from typing import List, Optional
import requests

import logging

from haystack import Document, component
from haystack.lazy_imports import LazyImport
from PIL import Image

logger = logging.getLogger(__name__)

with LazyImport(message="Run 'pip install transformers[torch,sentencepiece]'") as torch_and_transformers_import:
    import torch
    from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer, BlipProcessor, BlipForConditionalGeneration
    from PIL import Image

@component
class ImageCaptioner:
    def __init__(
        self,
        model_name: str = "Salesforce/blip-image-captioning-base",
    ):
        torch_and_transformers_import.check()
        self.model_name = model_name

        if model_name == "nlpconnect/vit-gpt2-image-captioning":
            self.model = VisionEncoderDecoderModel.from_pretrained(model_name)
            self.feature_extractor = ViTImageProcessor.from_pretrained(model_name)
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            max_length = 16
            num_beams = 4
            self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
        else:
            self.processor = BlipProcessor.from_pretrained(model_name)
            self.model = BlipForConditionalGeneration.from_pretrained(model_name)

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        
    @component.output_types(caption=str)
    def run(self, image_file_path: str) -> List[Document]:
        
        i_image = Image.open(image_file_path)
        if i_image.mode != "RGB":
            i_image = i_image.convert(mode="RGB")
        
        preds = []
        if self.model_name == "nlpconnect/vit-gpt2-image-captioning":
            pixel_values = self.feature_extractor(images=[i_image], return_tensors="pt").pixel_values
            pixel_values = pixel_values.to(self.device)

            output_ids = self.model.generate(pixel_values, **self.gen_kwargs)

            preds = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
            preds = [pred.strip() for pred in preds]
        else:

            inputs = self.processor([i_image], return_tensors="pt")
            output_ids = self.model.generate(**inputs)
            preds = self.processor.batch_decode(output_ids, skip_special_tokens=True)
            preds = [pred.strip() for pred in preds]
        
        # captions: List[Document] = []
        # for caption, image_file_path in zip(preds, image_file_paths):
        #     document = Document(content=caption, meta={"image_path": image_file_path})
        #     captions.append(document)
        return {"caption": preds[0]}