from __future__ import annotations

import asyncio
import base64
import io
import json
import os
import tempfile
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

import requests
from PIL import Image

from lagunavision.visual_pipeline import LagunaVisionImagePipeline, VisualProjectorSpec


class EndpointHandler:
    """Hugging Face Inference Endpoint handler for Laguna Vision checkpoints."""

    def __init__(self, path: str = "") -> None:
        self.model_dir = Path(path or ".").resolve()
        self.checkpoint_dir = self._resolve_checkpoint_dir()
        self.pipeline: LagunaVisionImagePipeline | None = None

    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
        payload = data.get("inputs", data)
        if not isinstance(payload, dict):
            raise ValueError("inputs must be an object with image and question fields")

        payload = _normalize_payload(payload)
        question = str(payload.get("question") or "").strip()
        if not question:
            raise ValueError("question is required")

        image_value = payload.get("image")
        if image_value is None:
            raise ValueError(
                "image is required as base64, a data URI, an HTTPS URL, a local path, "
                "or OpenAI-style messages content with image_url"
            )

        context = str(payload.get("context") or "")
        max_new_tokens = int(payload.get("max_new_tokens") or os.environ.get("LAGUNA_MAX_NEW_TOKENS", "128"))

        image = _load_image(image_value)
        with tempfile.NamedTemporaryFile(suffix=".png") as tmp:
            image.save(tmp.name)
            answer = asyncio.run(self._answer(Path(tmp.name), question, context, max_new_tokens))
        return {"answer": answer, "checkpoint": str(self.checkpoint_dir.relative_to(self.model_dir))}

    async def _answer(self, image: Path, question: str, context: str, max_new_tokens: int) -> str:
        if self.pipeline is None:
            self.pipeline = await self._load_pipeline()
        return await self.pipeline.answer_image(
            image=image,
            question=question,
            context=context,
            max_new_tokens=max_new_tokens,
        )

    async def _load_pipeline(self) -> LagunaVisionImagePipeline:
        spec_path = self.checkpoint_dir / "projector_spec.json"
        if not spec_path.exists():
            raise FileNotFoundError(f"missing projector_spec.json in {self.checkpoint_dir}")
        spec_row = json.loads(spec_path.read_text(encoding="utf-8"))
        spec = VisualProjectorSpec(
            input_dim=int(spec_row["input_dim"]),
            embedding_dim=int(spec_row["embedding_dim"]),
            hidden_dim=int(spec_row["hidden_dim"]),
            projector=spec_row.get("projector", "mlp"),
            visual_tokens=int(spec_row.get("visual_tokens", 64)),
            encoder=spec_row.get("encoder", "hf"),
            encoder_id=spec_row.get("encoder_id", ""),
            max_tiles=int(spec_row.get("max_tiles", 4)),
            patch_px=int(spec_row.get("patch_px", 32)),
        )
        lora_dir = self.checkpoint_dir / "lora"
        return await LagunaVisionImagePipeline.from_checkpoint(
            checkpoint=self.checkpoint_dir / "projector.pt",
            spec=spec,
            backbone_name=os.environ.get("LAGUNA_BACKBONE") or spec_row.get("backbone", "laguna"),
            model_id=os.environ.get("LAGUNA_MODEL_ID") or spec_row["model_id"],
            device=os.environ.get("LAGUNA_DEVICE", "auto"),
            vision_device=os.environ.get("LAGUNA_VISION_DEVICE", "auto"),
            lora_dir=lora_dir if lora_dir.exists() else None,
        )

    def _resolve_checkpoint_dir(self) -> Path:
        requested = os.environ.get("LAGUNA_CHECKPOINT_PATH", "latest").strip("/")
        candidates = [self.model_dir / requested, self.model_dir]
        for candidate in candidates:
            if (candidate / "projector.pt").exists() and (candidate / "projector_spec.json").exists():
                return candidate
        matches = sorted(self.model_dir.glob("**/projector.pt"), key=lambda path: path.stat().st_mtime, reverse=True)
        if matches:
            return matches[0].parent
        raise FileNotFoundError(
            f"no Laguna Vision checkpoint found under {self.model_dir}; expected latest/projector.pt"
        )


def _load_image(value: Any) -> Image.Image:
    if isinstance(value, bytes):
        return Image.open(io.BytesIO(value)).convert("RGB")
    if not isinstance(value, str):
        raise ValueError("image must be bytes or a string")

    if value.startswith("data:image/"):
        _, encoded = value.split(",", 1)
        return Image.open(io.BytesIO(base64.b64decode(encoded))).convert("RGB")

    parsed = urlparse(value)
    if parsed.scheme in {"http", "https"}:
        response = requests.get(value, timeout=20)
        response.raise_for_status()
        return Image.open(io.BytesIO(response.content)).convert("RGB")

    path = Path(value)
    if path.exists():
        return Image.open(path).convert("RGB")

    return Image.open(io.BytesIO(base64.b64decode(value))).convert("RGB")


def _normalize_payload(payload: dict[str, Any]) -> dict[str, Any]:
    if payload.get("messages") is None:
        return payload

    question_parts: list[str] = []
    image_value: Any = payload.get("image")
    for message in payload.get("messages") or []:
        if not isinstance(message, dict):
            continue
        content = message.get("content")
        if isinstance(content, str):
            question_parts.append(content)
            continue
        if not isinstance(content, list):
            continue
        for item in content:
            if not isinstance(item, dict):
                continue
            item_type = item.get("type")
            if item_type in {"text", "input_text"} and item.get("text"):
                question_parts.append(str(item["text"]))
            elif item_type in {"image_url", "input_image"}:
                image_url = item.get("image_url")
                if isinstance(image_url, dict):
                    image_value = image_url.get("url") or image_url.get("image_url") or image_value
                else:
                    image_value = image_url or item.get("url") or image_value

    normalized = dict(payload)
    normalized.setdefault("question", "\n".join(part.strip() for part in question_parts if part.strip()))
    if image_value is not None:
        normalized["image"] = image_value
    return normalized