"""
Custom inference handler for Hugging Face Inference Endpoints.

Serves the Marqo GCL ViT-L-14 open_clip checkpoint as a pure encoder:
given text and/or image(s), it returns L2-normalized embeddings, exactly
matching the notebook's encode_text / encode_image + normalize steps.

Request  (JSON):
    { "inputs": { "text":  ["a red dress", ...],          # optional
                  "image": ["<base64-encoded-image>", ...] } }  # optional

Response (JSON):
    { "text_embeddings":  [[... 768 floats ...], ...],    # if text sent
      "image_embeddings": [[... 768 floats ...], ...] }   # if image sent

Fusion (alpha*img + (1-alpha)*txt) is intentionally NOT done here — that
happens in the n8n Code node, so this endpoint stays a clean encoder.
"""

import os
import io
import glob
import base64

import torch
import open_clip
from PIL import Image


class EndpointHandler:
    def __init__(self, path: str = ""):
        # `path` is the local directory where the repo (including the .pt) is mounted.
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # Architecture of the checkpoint. ViT-L-14 for the Marqo GCL model.
        arch = os.environ.get("OPEN_CLIP_ARCH", "ViT-L-14")

        # Build the architecture WITHOUT pretrained weights, then load the GCL checkpoint.
        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
            arch, pretrained=None
        )
        self.tokenizer = open_clip.get_tokenizer(arch)

        ckpt_path = self._find_checkpoint(path)
        self._load_weights(ckpt_path)

        self.model = self.model.to(self.device).eval()

    # ---- locate the checkpoint file inside the repo ----
    def _find_checkpoint(self, path: str) -> str:
        # Allow an explicit override via env var (filename or absolute path).
        env = os.environ.get("CHECKPOINT_FILE")
        if env:
            return env if os.path.isabs(env) else os.path.join(path, env)

        candidates = (
            glob.glob(os.path.join(path, "*.pt"))
            + glob.glob(os.path.join(path, "*.bin"))
            + glob.glob(os.path.join(path, "*.safetensors"))
        )
        if not candidates:
            raise FileNotFoundError(
                f"No .pt/.bin/.safetensors checkpoint found in '{path}'. "
                "Set CHECKPOINT_FILE env var to the filename."
            )
        # If several exist, prefer the largest (the model weights, not a config shard).
        return max(candidates, key=os.path.getsize)

    # ---- load weights, tolerating the 'full_states' training-checkpoint layout ----
    def _load_weights(self, ckpt_path: str) -> None:
        # The Marqo file is named "...full_states.pt": a full training state that
        # wraps the model weights under a 'state_dict' key and may prefix keys with
        # 'module.' (from DDP training). open_clip.load_checkpoint normalizes both.
        try:
            open_clip.load_checkpoint(self.model, ckpt_path, strict=False)
            return
        except Exception as e:
            print(f"[handler] open_clip.load_checkpoint fallback triggered: {e}")

        # Manual fallback. weights_only=False is required because a full training
        # state contains non-tensor objects (optimizer state, epoch, etc.).
        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

        if isinstance(ckpt, dict):
            state_dict = ckpt.get("state_dict", ckpt.get("model", ckpt))
        else:
            state_dict = ckpt

        # Strip a leading 'module.' from DDP-saved keys.
        state_dict = {
            (k[len("module."):] if k.startswith("module.") else k): v
            for k, v in state_dict.items()
        }

        missing, unexpected = self.model.load_state_dict(state_dict, strict=False)
        if missing:
            print(f"[handler] missing keys (first 5): {list(missing)[:5]}")
        if unexpected:
            print(f"[handler] unexpected keys (first 5): {list(unexpected)[:5]}")

    @torch.no_grad()
    def __call__(self, data: dict) -> dict:
        # HF passes {"inputs": ..., "parameters": ...}; be lenient about shape.
        inputs = data.get("inputs", data) or {}
        if not isinstance(inputs, dict):
            return {"error": "Send an object under 'inputs' with 'text' and/or 'image'."}

        out = {}

        # ---- text embeddings ----
        texts = inputs.get("text")
        if texts:
            if isinstance(texts, str):
                texts = [texts]
            tokens = self.tokenizer(texts).to(self.device)
            feats = self.model.encode_text(tokens)
            feats = feats / feats.norm(dim=-1, keepdim=True)
            out["text_embeddings"] = feats.cpu().tolist()

        # ---- image embeddings (base64-encoded bytes) ----
        images = inputs.get("image")
        if images:
            if isinstance(images, str):
                images = [images]
            tensors = []
            for b64 in images:
                raw = base64.b64decode(b64)
                img = Image.open(io.BytesIO(raw)).convert("RGB")
                tensors.append(self.preprocess(img))
            batch = torch.stack(tensors).to(self.device)
            feats = self.model.encode_image(batch)
            feats = feats / feats.norm(dim=-1, keepdim=True)
            out["image_embeddings"] = feats.cpu().tolist()

        if not out:
            return {
                "error": "Provide 'text' and/or 'image' (base64) under 'inputs'.",
                "example": {"inputs": {"text": ["a red dress"]}},
            }
        return out