import os
import requests
import tempfile
from PIL import Image
from tools.interfaces.interface_image_analyzer import ImageAnalysisInterface
from transformers import BlipProcessor, BlipForConditionalGeneration
import torch

class BlipImageAnalyzer(ImageAnalysisInterface):
    def __init__(self, huggingface_token: str, base_url: str):
        self.huggingface_token = huggingface_token
        self.base_url = base_url
        self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
        self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

    def analyze(self, image_filename: str) -> str:
        tmp_path = None
        try:
            # Vérifier si le fichier est local
            if os.path.isfile(image_filename):
                # Chemin local
                image_path = image_filename
            else:
                # Sinon, on suppose que c'est un fichier distant à télécharger
                headers = {"Authorization": f"Bearer {self.huggingface_token}"}
                url = f"{self.base_url}{image_filename}"
                response = requests.get(url, headers=headers)
                response.raise_for_status()

                # Enregistrer temporairement l'image
                suffix = os.path.splitext(image_filename)[-1]
                with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
                    tmp.write(response.content)
                    tmp_path = tmp.name
                image_path = tmp_path

            # Ouvrir et traiter l'image
            image = Image.open(image_path).convert("RGB")
            inputs = self.processor(images=image, return_tensors="pt")

            # Générer la légende
            with torch.no_grad():
                output = self.model.generate(**inputs)
            caption = self.processor.decode(output[0], skip_special_tokens=True)

            return f"(Vision-BLIP) Description de l'image : {caption}"

        except Exception as e:
            return f"(Vision-BLIP) Erreur d’analyse : {e}"

        finally:
            # Supprimer le fichier temporaire si nécessaire
            if tmp_path and os.path.isfile(tmp_path):
                os.remove(tmp_path)