import base64
import io
from typing import Any, Dict, Optional

import numpy as np
import torch
from PIL import Image
from diffusers import MarigoldIntrinsicsPipeline, utils as diffusers_utils

# We will import MarigoldNormalsPipeline lazily to avoid loading it unless asked
NormalsPipe = None


def _clean_b64(s: str) -> bytes:
    s = s.strip()
    if s.startswith("data:image"):
        s = s.split(",", 1)[1]
    s = "".join(s.split())  # strip whitespace/newlines
    pad = (-len(s)) % 4
    if pad:
        s += "=" * pad
    try:
        return base64.b64decode(s, validate=False)
    except Exception:
        return base64.urlsafe_b64decode(s)


def _image_from_inputs(obj: Any) -> Image.Image:
    if isinstance(obj, dict):
        if "image_url" in obj:
            return diffusers_utils.load_image(obj["image_url"]).convert("RGB")
        if "image_b64" in obj:
            raw = _clean_b64(obj["image_b64"])
            return Image.open(io.BytesIO(raw)).convert("RGB")
        raise ValueError("inputs dict must contain 'image_url' or 'image_b64'.")

    if isinstance(obj, (str, bytes, bytearray)):
        if isinstance(obj, (bytes, bytearray)):
            obj = obj.decode("utf-8", errors="ignore")
        s = obj.strip()
        if s.startswith("http://") or s.startswith("https://"):
            return diffusers_utils.load_image(s).convert("RGB")
        raw = _clean_b64(s)
        return Image.open(io.BytesIO(raw)).convert("RGB")

    raise ValueError("Unsupported 'inputs' type. Provide URL string, base64 string, or {'image_url'|'image_b64'}.")


def _to_uint8_png_b64(img_like: Any) -> Optional[str]:
    """
    Best-effort: accept PIL.Image or np.ndarray / torch.Tensor (H,W,3) or (3,H,W),
    map to [0,255] and return PNG b64 string.
    """
    try:
        if isinstance(img_like, Image.Image):
            buf = io.BytesIO()
            img_like.save(buf, format="PNG")
            return base64.b64encode(buf.getvalue()).decode("utf-8")

        if isinstance(img_like, torch.Tensor):
            arr = img_like.detach().cpu().float().numpy()
        elif isinstance(img_like, np.ndarray):
            arr = img_like
        else:
            return None

        # Move channels-last if needed
        if arr.ndim == 3 and arr.shape[0] in (1, 3) and arr.shape[-1] not in (1, 3):
            arr = np.transpose(arr, (1, 2, 0))
        if arr.ndim != 3 or arr.shape[-1] not in (1, 3):
            # If it's something else (e.g., batch), try take first item
            if arr.ndim == 4:  # (B,H,W,C)
                arr = arr[0]
            else:
                return None

        # Heuristic normalization
        arr = arr.astype(np.float32)
        a_min, a_max = float(np.nanmin(arr)), float(np.nanmax(arr))
        if a_max <= 1.0 and a_min >= 0.0:
            # already [0,1]
            pass
        elif a_max <= 1.0 and a_min >= -1.0:
            # [-1,1] normals style
            arr = (arr + 1.0) * 0.5
        else:
            # generic min-max
            rng = (a_max - a_min) if (a_max - a_min) > 1e-6 else 1.0
            arr = (arr - a_min) / rng

        arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
        img = Image.fromarray(arr)
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return base64.b64encode(buf.getvalue()).decode("utf-8")
    except Exception:
        return None


class EndpointHandler:
    """HF Inference Toolkit-compatible class handler with Intrinsics (as-is) + optional Normals."""

    def __init__(self, path: Optional[str] = None):
        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if device == "cuda" else torch.float32

        # ---- Intrinsics pipeline (UNCHANGED from your working code) ----
        self.pipe = MarigoldIntrinsicsPipeline.from_pretrained(
            "prs-eth/marigold-iid-lighting-v1-1",
            torch_dtype=dtype,
        ).to(device)
        try:
            self.pipe.set_progress_bar_config(disable=True)
        except Exception:
            pass

        # Expected targets for lighting
        self.targets = ["albedo", "shading", "residual"]

        # ---- Normals (lazy) ----
        self._normals_pipe = None
        self._device = device
        self._dtype = dtype

        # Defaults
        self.default_intr_steps = 4  # as before
        self.fixed_processing_res = 768  # per your latest preference

    # ---------------- Normals helpers ----------------
    def _ensure_normals_loaded(self):
        global NormalsPipe
        if self._normals_pipe is not None:
            return
        if NormalsPipe is None:
            # Lazy import to avoid loading unless requested
            from diffusers.pipelines.marigold.pipeline_marigold_normals import MarigoldNormalsPipeline as _NP
            NormalsPipe = _NP
        self._normals_pipe = NormalsPipe.from_pretrained(
            "prs-eth/marigold-normals-v1-1",
            torch_dtype=self._dtype,
        ).to(self._device)
        try:
            self._normals_pipe.set_progress_bar_config(disable=True)
        except Exception:
            pass

    def _run_normals(self, pil_img: Image.Image) -> Dict[str, Any]:
        try:
            self._ensure_normals_loaded()
            # Do NOT pass processing_res to avoid the earlier error.
            # Most versions work with just the image (and optionally steps if supported).
            out = self._normals_pipe(pil_img)

            # Try common fields in order
            pred = None
            for key in ("prediction", "pred_normals"):
                if hasattr(out, key):
                    pred = getattr(out, key)
                    break
            if pred is None and hasattr(out, "images") and out.images:
                # Some pipelines return PIL in images[0]
                pred = out.images[0]

            # Convert to info + small preview b64
            if isinstance(pred, torch.Tensor):
                arr = pred.detach().cpu().numpy()
                shape = list(arr.shape)
                dtype_str = str(arr.dtype)
                png_b64 = _to_uint8_png_b64(arr)
            elif isinstance(pred, np.ndarray):
                shape = list(pred.shape)
                dtype_str = str(pred.dtype)
                png_b64 = _to_uint8_png_b64(pred)
            elif isinstance(pred, Image.Image):
                shape = [pred.height, pred.width, 3]
                dtype_str = "uint8"
                png_b64 = _to_uint8_png_b64(pred)
            else:
                return {"error": f"Unexpected prediction type: {type(pred)}"}

            if png_b64 is None:
                return {"error": "Too many dimensions: 4 > 3."}

            return {
                "tensors_info": {"shape": shape, "dtype": dtype_str},
                "preview_png_b64": png_b64,
            }
        except torch.cuda.OutOfMemoryError:
            return {"error": "GPU out of memory. Try a smaller image."}
        except Exception as e:
            return {"error": f"Normals prediction failed: {str(e)}"}

    # ---------------- Intrinsics (unchanged behavior) ----------------
    def _run_intrinsics(self, pil_img: Image.Image, num_steps: int) -> Dict[str, Any]:
        # Inference (as before)
        with torch.inference_mode():
            out = self.pipe(pil_img, num_inference_steps=num_steps)

        pred = out.prediction
        # Support torch.Tensor OR numpy.ndarray
        if isinstance(pred, torch.Tensor):
            pred_np = pred.detach().cpu().numpy()
        elif isinstance(pred, np.ndarray):
            pred_np = pred
        else:
            return {"error": f"Unexpected prediction type: {type(pred)}"}

        # pred_np expected shape: (T, H, W, 3) or (H, W, 3) depending on pipeline
        shape = list(pred_np.shape)
        dtype_str = str(pred_np.dtype)

        # Visualizations (best-effort) — unchanged
        vis = {}
        vis_err = None
        try:
            # visualize_intrinsics expects numpy prediction and the pipeline's target_properties
            vis_list = self.pipe.image_processor.visualize_intrinsics(pred_np, self.pipe.target_properties)
            # vis_list is a list over batch; we run single image, so index 0
            if vis_list and isinstance(vis_list[0], dict):
                for name, im in vis_list[0].items():
                    if isinstance(im, Image.Image):
                        buf = io.BytesIO()
                        im.save(buf, format="PNG")
                        vis[name] = base64.b64encode(buf.getvalue()).decode("utf-8")
        except Exception as e:
            vis_err = f"Failed to visualize: {e}"

        return {
            "targets": self.targets,
            "tensors_info": {"shape": shape, "dtype": dtype_str},
            "uncertainty_available": False,
            "visualizations_png_b64": vis or None,
            "visualizations_error": vis_err,
        }

    # ---------------- Entry point ----------------
    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        try:
            if "inputs" not in data:
                return {"error": "Body must include an 'inputs' key."}

            # Parameters can be at top-level or under 'parameters'
            params = {}
            params.update({k: v for k, v in data.items() if k != "inputs"})
            if isinstance(data.get("parameters"), dict):
                params.update(data["parameters"])

            # Which outputs to return?
            ret = params.get("return", ["intrinsics"])
            if isinstance(ret, str):
                ret = [ret]
            ret = [r.lower() for r in ret]

            # Steps (intrinsics)
            try:
                num_steps = int(params.get("num_inference_steps", self.default_intr_steps))
            except Exception:
                num_steps = self.default_intr_steps

            # Decode image
            try:
                pil_img = _image_from_inputs(data["inputs"])
            except Exception as e:
                return {"error": f"Failed to decode 'inputs': {e}"}

            result: Dict[str, Any] = {"outputs": {}}

            if "normals" in ret:
                result["outputs"]["normals"] = self._run_normals(pil_img)

            if "intrinsics" in ret or not ret:
                intr_res = self._run_intrinsics(pil_img, num_steps=num_steps)
                # sanity check that vis contains expected keys when available
                if isinstance(intr_res, dict) and "targets" in intr_res:
                    # Keep your original “intrinsics-only” structure nested
                    result["outputs"]["intrinsics"] = intr_res
                else:
                    result["outputs"]["intrinsics"] = {
                        "error": "Intrinsics pipeline did not return expected outputs (albedo, diffuse_shading, residual)."
                    }

            # If only one branch requested, return only that branch in 'outputs'
            return result

        except Exception as e:
            return {"error": f"{e}"}
        finally:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()