from typing import Dict, Any
import requests
import io
import base64
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity


class EndpointHandler:
    def __init__(self, path=""):
        self.processor = CLIPProcessor.from_pretrained(path)
        self.model = CLIPModel.from_pretrained(path)

    def __call__(self, data: Dict[str, Any]) -> Dict:
        print("this shows the custom endpoint handler is being called")
        inputs = data.pop("inputs", data)
        text = inputs.pop("text")
        if "image_url" in inputs:
            image_url = inputs.pop("image_url")
            image = Image.open(requests.get(image_url, stream=True).raw)
        else:
            image = inputs.pop("image")
            image = Image.open(io.BytesIO(base64.b64decode(image)))
        processed_inputs = self.processor(text=text, images=image,
                                          return_tensors="pt", padding=True, truncation=True)
        outputs = self.model(**processed_inputs)
        embedding_similarity = cosine_similarity(outputs.text_embeds.detach().numpy(),
                                                 outputs.image_embeds.detach().numpy())[0][0].item()
        return {"text_embedding": outputs.text_embeds[0].tolist(),
                "image_embedding": outputs.image_embeds[0].tolist(),
                "embedding_similarity": embedding_similarity}