import torch
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import numpy as np

class CLIPEmbedder:
    def __init__(self, model_name: str = "openai/clip-vit-base-patch32"):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = CLIPModel.from_pretrained(model_name).to(self.device)
        self.processor = CLIPProcessor.from_pretrained(model_name)

    def get_text_embedding(self, text: str) -> np.ndarray:
        inputs = self.processor(text=[text], return_tensors="pt", padding=True).to(self.device)
        with torch.no_grad():
            text_features = self.model.get_text_features(**inputs)
        if not isinstance(text_features, torch.Tensor):
            if hasattr(text_features, "pooler_output"):
                text_features = text_features.pooler_output
            elif isinstance(text_features, tuple):
                text_features = text_features[0]
        text_features = text_features / text_features.norm(p=2, dim=-1, keepdim=True)
        return text_features.cpu().numpy()[0]

    def get_image_embedding(self, image: Image.Image) -> np.ndarray:
        if image.mode != "RGB":
            image = image.convert("RGB")
        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
        with torch.no_grad():
            image_features = self.model.get_image_features(**inputs)
        if not isinstance(image_features, torch.Tensor):
            if hasattr(image_features, "pooler_output"):
                image_features = image_features.pooler_output
            elif isinstance(image_features, tuple):
                image_features = image_features[0]
        image_features = image_features / image_features.norm(p=2, dim=-1, keepdim=True)
        return image_features.cpu().numpy()[0]