import gradio as gr
import spaces
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer
from transformers.image_utils import load_image
from threading import Thread
import time
import torch
from PIL import Image
import uuid
import io
import os

# Fine-tuned for OCR-based tasks from Qwen's [ Qwen/Qwen2-VL-2B-Instruct ]
MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = Qwen2VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.float16
).to("cuda").eval()

# Supported media extensions
image_extensions = Image.registered_extensions()
video_extensions = ("avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg", "wav", "gif", "webm", "m4v", "3gp")

def identify_and_save_blob(blob_path):
    """Identifies if the blob is an image or video and saves it accordingly."""
    try:
        with open(blob_path, 'rb') as file:
            blob_content = file.read()
            
            # Try to identify if it's an image
            try:
                Image.open(io.BytesIO(blob_content)).verify()  # Check if it's a valid image
                extension = ".png"  # Default to PNG for saving
                media_type = "image"
            except (IOError, SyntaxError):
                # If it's not a valid image, assume it's a video
                extension = ".mp4"  # Default to MP4 for saving
                media_type = "video"
            
            # Create a unique filename
            filename = f"temp_{uuid.uuid4()}_media{extension}"
            with open(filename, "wb") as f:
                f.write(blob_content)
                
            return filename, media_type
            
    except FileNotFoundError:
        raise ValueError(f"The file {blob_path} was not found.")
    except Exception as e:
        raise ValueError(f"An error occurred while processing the file: {e}")

def process_vision_info(messages):
    """Processes vision inputs (images and videos) from messages."""
    image_inputs = []
    video_inputs = []
    for message in messages:
        for content in message["content"]:
            if content["type"] == "image":
                image_inputs.append(load_image(content["image"]))
            elif content["type"] == "video":
                video_inputs.append(content["video"])
    return image_inputs, video_inputs

@spaces.GPU
def model_inference(input_dict, history):
    text = input_dict["text"]
    files = input_dict["files"]

    # Process media files (images or videos)
    media_paths = []
    media_types = []
    for file in files:
        if file.endswith(tuple([i for i, f in image_extensions.items()])):
            media_type = "image"
        elif file.endswith(video_extensions):
            media_type = "video"
        else:
            try:
                file, media_type = identify_and_save_blob(file)
            except Exception as e:
                gr.Error(f"Unsupported media type: {e}")
                return
        media_paths.append(file)
        media_types.append(media_type)

    # Validate input
    if text == "" and not media_paths:
        gr.Error("Please input a query and optionally image(s) or video(s).")
        return
    if text == "" and media_paths:
        gr.Error("Please input a text query along with the image(s) or video(s).")
        return

    # Prepare messages for the model
    messages = [
        {
            "role": "user",
            "content": [
                *[{"type": media_type, media_type: media_path} for media_path, media_type in zip(media_paths, media_types)],
                {"type": "text", "text": text},
            ],
        }
    ]

    # Apply chat template and process inputs
    prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # Process vision inputs (images and videos)
    image_inputs, video_inputs = process_vision_info(messages)

    # Ensure video_inputs is not empty
    if not video_inputs:
        video_inputs = None

    inputs = processor(
        text=[prompt],
        images=image_inputs if image_inputs else None,
        videos=video_inputs if video_inputs else None,
        return_tensors="pt",
        padding=True,
    ).to("cuda")

    # Set up streamer for real-time output
    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)

    # Start generation in a separate thread
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    # Stream the output
    buffer = ""
    yield "Thinking..."
    for new_text in streamer:
        buffer += new_text
        # Remove <|im_end|> or similar tokens from the output
        buffer = buffer.replace("<|im_end|>", "")
        time.sleep(0.01)
        yield buffer

# Example inputs
examples = [
    [{"text": "Describe the video.", "files": ["examples/demo.mp4"]}],
    [{"text": "Extract JSON from the image", "files": ["example_images/document.jpg"]}],
    [{"text": "summarize the letter", "files": ["examples/1.png"]}],
    [{"text": "Describe the photo", "files": ["examples/3.png"]}],
    [{"text": "Extract as JSON table from the table", "files": ["examples/4.jpg"]}],
    [{"text": "Summarize the full image in detail", "files": ["examples/2.jpg"]}],
    [{"text": "Describe this image.", "files": ["example_images/campeones.jpg"]}],
    [{"text": "What is this UI about?", "files": ["example_images/s2w_example.png"]}],
    [{"text": "Can you describe this image?", "files": ["example_images/newyork.jpg"]}],
    [{"text": "Can you describe this image?", "files": ["example_images/dogs.jpg"]}],
    [{"text": "Where do the severe droughts happen according to this diagram?", "files": ["example_images/examples_weather_events.png"]}],
]

demo = gr.ChatInterface(
    fn=model_inference,
    description="# **Multimodal OCR**",
    examples=examples,
    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image", "video"], file_count="multiple"),
    stop_btn="Stop Generation",
    multimodal=True,
    cache_examples=False,
)

demo.launch(debug=True, share=True)