from PIL import Image
from kaggle_gpu_server.engine.plugin_base import ModelPlugin, PluginCapability, ModelCategory

class DepthAnythingV2Plugin(ModelPlugin):
    name = "depth_anything_v2"
    model_id = "depth-anything/Depth-Anything-V2-Small-hf"
    capability = PluginCapability.DEPTH_ESTIMATION
    category = ModelCategory.LIGHTWEIGHT
    vram_estimate_mb = 400
    version = "1.0.0"
    description = "Estimates depth using Depth-Anything-V2-Small"

    def __init__(self):
        super().__init__()
        self.pipe = None

    def load(self) -> bool:
        if self._loaded:
            return True
        try:
            from transformers import pipeline
            self.pipe = pipeline("depth-estimation", model=self.model_id, device=self._device)
            self._loaded = True
            return True
        except Exception as e:
            print(f"❌ Failed to load Depth Anything: {e}")
            return False

    def unload(self) -> None:
        if not self._loaded:
            return
        self.pipe = None
        self._loaded = False
        import torch
        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def _execute(self, inputs: dict) -> dict:
        image = inputs.get("image")
        if image is None:
            raise ValueError("Input 'image' is required")
            
        if self.pipe is None:
            # Fallback empty depth map
            fallback = Image.new("L", image.size, 128)
            return {"depth_map": fallback, "image": image}
            
        result = self.pipe(image)
        depth_map = result["depth"]
        
        # Ensure depth map is resized to original image dimensions if transformers pipeline modified it
        if depth_map.size != image.size:
            depth_map = depth_map.resize(image.size, Image.Resampling.BILINEAR)
            
        return {
            "depth_map": depth_map,
            "image": image
        }