from typing import Dict, Any, List, Union
import base64
from io import BytesIO

import torch
import numpy as np
from PIL import Image
from transformers import XCLIPProcessor, XCLIPModel


class EndpointHandler:
    def __init__(self, path=""):
        self.model = XCLIPModel.from_pretrained(path)
        self.processor = XCLIPProcessor.from_pretrained(path)
        self.model.eval()
        try:
            self.num_frames = self.model.config.vision_config.num_frames
        except AttributeError:
            self.num_frames = 8

    def _decode(self, b64: str) -> Image.Image:
        return Image.open(BytesIO(base64.b64decode(b64))).convert("RGB")

    def _sample(self, frames: List[Image.Image]) -> List[np.ndarray]:
        idx = np.linspace(0, len(frames) - 1, self.num_frames).round().astype(int)
        return [np.array(frames[i]) for i in idx]

    def __call__(self, data: Dict[str, Any]) -> Union[Dict, List]:
        inputs = data.get("inputs", data)
        frames_b64 = inputs["frames"] if isinstance(inputs, dict) else inputs
        labels = inputs.get("candidate_labels") if isinstance(inputs, dict) else None

        frames = [self._decode(f) for f in frames_b64]
        if not frames:
            raise ValueError("No frames provided")
        video = self._sample(frames)

        if labels:
            proc = self.processor(text=labels, videos=video,
                                  return_tensors="pt", padding=True)
            with torch.no_grad():
                probs = self.model(**proc).logits_per_video.softmax(dim=1)[0]
            return [{"label": l, "score": float(s)} for l, s in zip(labels, probs)]

        proc = self.processor(videos=video, return_tensors="pt")
        with torch.no_grad():
            feats = self.model.get_video_features(pixel_values=proc["pixel_values"])
        feats = feats / feats.norm(p=2, dim=-1, keepdim=True)
        return {"embedding": feats[0].tolist(), "dim": int(feats.shape[-1])}