"""
models/depth_estimation.py
──────────────────────────
TransformersDepthLoader  — wraps any HuggingFace depth-estimation
model that follows the transformers pipeline API (Depth Anything v2,
DPT-Large, MiDaS, ZoeDepth, etc.).

Inputs  (run kwargs)
──────────────────────────────────────────────────────────────
  image   : PIL.Image  (RGB)  — the image to estimate depth for

Outputs (returned dict)
──────────────────────────────────────────────────────────────
  depth_raw        : np.ndarray  (H×W float32, unnormalised)
  depth_normalised : np.ndarray  (H×W float32 in [0,1])
  depth_colourmap  : PIL.Image   (false-colour for display)
  depth_uint16     : PIL.Image   (16-bit greyscale PNG for export)
  model            : str
"""

from __future__ import annotations

import logging
import os
from typing import Any

import numpy as np
from PIL import Image

from models.base_loader import BaseLoader
from utils.image_utils import normalise_depth, depth_to_colormap, depth_to_uint16, resize_image

logger = logging.getLogger(__name__)

def _hf_token() -> str | None:
    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None


class TransformersDepthLoader(BaseLoader):

    def load(self) -> None:
        if self._loaded:
            return
        logger.info("Loading depth model: %s", self.model_id)

        from transformers import pipeline as hf_pipeline

        token = _hf_token()
        pipeline_kwargs: dict = {
            "task": "depth-estimation",
            "model": self.model_id,
            "device": self.device,
        }
        if token:
            pipeline_kwargs["token"] = token

        try:
            self.pipe = hf_pipeline(**pipeline_kwargs)
        except (OSError, EnvironmentError) as hub_err:
            err_str = str(hub_err).lower()
            if any(k in err_str for k in ("not cached", "fetch metadata", "connection", "network",
                                           "offline", "cannot reach", "name or service not known")):
                logger.warning(
                    "Hub unreachable for depth model %s. Retrying with local_files_only=True …",
                    self.model_id,
                )
                try:
                    self.pipe = hf_pipeline(**{**pipeline_kwargs, "model": self.model_id},
                                            local_files_only=True)
                    logger.info("Loaded depth model %s from local cache.", self.model_id)
                except Exception as cache_err:
                    raise OSError(
                        f"Cannot load depth model {self.model_id}: Hub unreachable and no local cache.\n"
                        f"  Hub error  : {hub_err}\n"
                        f"  Cache error: {cache_err}\n"
                        "Tip: ensure internet access on first run, or set HF_TOKEN for gated models."
                    ) from cache_err
            else:
                raise

        self._loaded = True
        logger.info("Depth model ready: %s", self.model_id)

    def run(self, **inputs: Any) -> dict[str, Any]:
        if not self._loaded:
            self.load()

        image: Image.Image = inputs["image"]
        image_size: int = int(self.kwargs.get("image_size", 518))

        # Resize to the model's preferred resolution
        image = resize_image(image, max_side=image_size).convert("RGB")

        logger.info("Estimating depth  image_size=%d  model=%s", image_size, self.model_id)

        result = self.pipe(image)

        # HF depth-estimation pipeline returns {"depth": PIL.Image, "predicted_depth": tensor}
        depth_tensor = result.get("predicted_depth")
        if depth_tensor is not None:
            import torch
            depth_raw = depth_tensor.squeeze().detach().cpu().numpy().astype(np.float32)
        else:
            # Fallback: use the greyscale PIL image
            depth_raw = np.array(result["depth"]).astype(np.float32)

        depth_norm = normalise_depth(depth_raw)
        depth_colour = depth_to_colormap(depth_raw)
        depth_16 = depth_to_uint16(depth_raw)

        return {
            "depth_raw": depth_raw,
            "depth_normalised": depth_norm,
            "depth_colourmap": depth_colour,
            "depth_uint16": depth_16,
            "model": self.model_id,
        }