from typing import List, Union
from PIL import Image, ImageDraw
import torch

from diffusers.modular_pipelines import (
    PipelineState,
    ModularPipelineBlocks,
    InputParam,
    ComponentSpec,
    OutputParam,
)
from transformers import AutoProcessor, AutoModelForCausalLM


class Florence2ImageAnnotatorBlock(ModularPipelineBlocks):
    @property
    def expected_components(self):
        return [
            ComponentSpec(
                name="image_annotator",
                type_hint=AutoModelForCausalLM,
                repo="microsoft/Florence-2-large",
            ),
            ComponentSpec(
                name="image_annotator_processor",
                type_hint=AutoProcessor,
                repo="microsoft/Florence-2-large",
            ),
        ]

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam(
                "image",
                Image,
                required=True,
                description="Image(s) to annotate",
            ),
            InputParam(
                "annotation_task_prompt",
                Union[str, List[str]],
                required=True,
                description="""Annotation Task to perform on the image.
                """,
            ),
        ]

    @property
    def intermediates_outputs(self) -> List[OutputParam]:
        return [
            OutputParam(
                "mask",
                type_hint=torch.Tensor,
                description="Depth Map(s) of input Image(s)",
            ),
        ]

    def annotate_image(self, image, prompt):
        inputs = self.image_annotator_processor(
            text=prompt, images=image, return_tensors="pt"
        )
        generated_ids = self.annotator.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs["pixel_values"],
            max_new_tokens=1024,
            early_stopping=False,
            do_sample=False,
            num_beams=3,
        )
        annotations = self.image_annotator_processor.batch_decode(
            generated_ids, skip_special_tokens=False
        )[0]
        annotations = self.image_annotator_processor.post_process_generation(
            annotations, task=prompt, image_size=(image.height, image.width)
        )

        return annotations

    def prepare_mask(self, images, annotations):
        masks = []
        for image, annotation in zip(images, annotations):
            mask_image = Image.new("L", image.size, 0)
            draw = ImageDraw.Draw(mask_image)
            draw.polygon(annotation["polygon"], fill="white")
            masks.append(mask_image)

        return masks

    @torch.no_grad()
    def __call__(self, pipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        images = block_state.image
        annotation_task_prompt = block_state.annotation_task_prompt

        if not isinstance(annotation_task_prompt, list):
            annotation_task_prompt = [annotation_task_prompt]

        if len(images) != len(annotation_task_prompt):
            raise ValueError("Number of images and annotation prompts must match")

        annotations = self.annotate_image(images, annotation_task_prompt)
        block_state.mask = self.prepare_mask(images, annotations)

        self.set_block_state(block_state)