import base64
import io
import os
from typing import Any, Dict, List
from urllib.parse import urlparse
from urllib.request import urlopen

import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = AutoModel.from_pretrained(path, trust_remote_code=True)
        self.processor = AutoProcessor.from_pretrained(path, trust_remote_code=True)
        self.model.to(self.device)
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        payload = data.pop("inputs", data)
        parameters = data.pop("parameters", {}) or {}

        texts = self._coerce_texts(payload)
        images = self._coerce_images(payload)
        if not texts and not images:
            raise ValueError(
                "Expected `inputs` to include `text`/`texts` and/or `image`/`images`."
            )

        result: Dict[str, Any] = {}
        with torch.no_grad():
            text_embeds = None
            image_embeds = None

            if texts:
                text_inputs = self.processor(text=texts, return_tensors="pt")
                text_inputs = self._move_to_device(text_inputs)
                text_embeds = self.model(**text_inputs).text_embeds
                result["text_embedding"] = text_embeds.cpu().tolist()

            if images:
                image_inputs = self.processor(images=images, return_tensors="pt")
                image_inputs = self._move_to_device(image_inputs)
                image_embeds = self.model(**image_inputs).image_embeds
                result["image_embedding"] = image_embeds.cpu().tolist()

            if text_embeds is not None and image_embeds is not None:
                scores = image_embeds @ text_embeds.t()
                result["scores"] = scores.cpu().tolist()
                if parameters.get("return_probs", True):
                    result["probs"] = scores.softmax(dim=-1).cpu().tolist()
                if parameters.get("return_logits", False):
                    logit_scale = self.model.model.logit_scale.exp()
                    result["logits_per_image"] = (
                        (logit_scale * image_embeds @ text_embeds.t()).cpu().tolist()
                    )

        return result

    def _move_to_device(self, batch: Dict[str, Any]) -> Dict[str, Any]:
        moved = {}
        for key, value in batch.items():
            moved[key] = value.to(self.device) if hasattr(value, "to") else value
        return moved

    def _coerce_texts(self, payload: Any) -> List[str]:
        if isinstance(payload, str):
            return [payload]
        if not isinstance(payload, dict):
            return []

        texts = payload.get("text", payload.get("texts"))
        if texts is None:
            return []
        if isinstance(texts, str):
            return [texts]
        return [str(item) for item in texts]

    def _coerce_images(self, payload: Any) -> List[Image.Image]:
        if not isinstance(payload, dict):
            return []

        images = payload.get("image", payload.get("images"))
        if images is None:
            return []
        if not isinstance(images, (list, tuple)):
            images = [images]
        return [self._load_image(item) for item in images]

    def _load_image(self, value: Any) -> Image.Image:
        if isinstance(value, Image.Image):
            return value.convert("RGB")

        if isinstance(value, dict):
            for key in ("data", "image", "url", "path"):
                if key in value:
                    value = value[key]
                    break

        if not isinstance(value, str):
            raise TypeError(f"Unsupported image input type: {type(value)!r}")

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

        parsed = urlparse(value)
        if parsed.scheme in ("http", "https"):
            with urlopen(value) as response:
                return Image.open(io.BytesIO(response.read())).convert("RGB")

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

        try:
            return Image.open(io.BytesIO(base64.b64decode(value))).convert("RGB")
        except Exception as exc:
            raise ValueError("Unsupported image string. Use URL, local path, or base64.") from exc