# Importing necessary libraries
import sys
from typing import Optional
from unittest.mock import patch
from PIL import Image
import torch
from transformers import AutoProcessor, AutoModelForCausalLM
from transformers.dynamic_module_utils import get_imports
import spaces
import gradio as gr

# Local imports
from src.logger import logging
from src.exception import CustomExceptionHandling


# Florence-2's custom modeling file (modeling_florence2.py) unconditionally
# lists "flash_attn" as a required import, even though we load the model
# with attn_implementation="eager" (which never touches flash_attn).
# Actually installing a working flash-attn build inside a Space is slow
# and often fails outright, since it needs a prebuilt wheel that exactly
# matches the image's torch/CUDA/Python ABI -- if no match exists, pip
# silently produces a non-functional stub package, which still leaves
# the model unable to load. Since flash_attn isn't actually needed here,
# the standard fix is to patch transformers' import-checker to drop
# "flash_attn" from the required-imports list, only for this one file.
# Reference: https://huggingface.co/microsoft/Florence-2-large-ft/discussions/4
def _fixed_get_imports(filename):
    imports = get_imports(filename)
    if str(filename).endswith("/modeling_florence2.py") and "flash_attn" in imports:
        imports.remove("flash_attn")
    return imports


# Load model and processor from Hugging Face
model_id = "microsoft/Florence-2-large-ft"
try:
    with patch("transformers.dynamic_module_utils.get_imports", _fixed_get_imports):
        model = (
            AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="eager", torch_dtype=torch.float16, trust_remote_code=True)
            .to("cuda")
            .eval()
        )
        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    logging.info("Model and processor loaded successfully.")

# Handle exceptions that may occur during the process
except Exception as e:
    # Custom exception handling
    raise CustomExceptionHandling(e, sys) from e


@spaces.GPU
def run_example(
    task_prompt: str, image: Image.Image, text_input: Optional[str] = None
) -> str:
    """
    Runs an example using the given task prompt and image.

    Args:
        - task_prompt (str): The task prompt for the example.
        - image (PIL.Image.Image): The image to be processed.
        - text_input (str, optional): Additional text input to be appended to the task prompt. Defaults to None.

    Returns:
        str: The parsed answer generated by the model.
    """
    try:
        # Check if image is None
        if image is None:
            gr.Warning("Please provide an image.")

        # If there is no text input, use the task prompt as the prompt
        prompt = task_prompt if text_input is None else task_prompt + text_input

        # Process the image and text input
        inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda", torch.float16)

        # Generate the answer using the model
        generated_ids = model.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs["pixel_values"],
            max_new_tokens=1024,
            early_stopping=False,
            do_sample=False,
            num_beams=3,
        )
        generated_text = processor.batch_decode(
            generated_ids, skip_special_tokens=False
        )[0]
        parsed_answer = processor.post_process_generation(
            generated_text, task=task_prompt, image_size=(image.width, image.height)
        )

        # Return the parsed answer
        return parsed_answer

    # Handle exceptions that may occur during the process
    except Exception as e:
        # Custom exception handling
        raise CustomExceptionHandling(e, sys) from e