import gradio as gr
from transformers import (
    AutoProcessor,
    Qwen2_5_VLForConditionalGeneration,
    TextIteratorStreamer,
    AutoModelForCausalLM,
    AutoTokenizer,
)
from transformers.image_utils import load_image
from threading import Thread
import time
import torch
import spaces
import cv2
import numpy as np
from PIL import Image

# -----------------------
# Progress Bar Helper
# -----------------------
def progress_bar_html(label: str) -> str:
    """
    Returns an HTML snippet for a thin progress bar with a label.
    The progress bar is styled as a dark animated bar.
    """
    return f'''
<div style="display: flex; align-items: center;">
    <span style="margin-right: 10px; font-size: 14px;">{label}</span>
    <div style="width: 110px; height: 5px; background-color: #9370DB; border-radius: 2px; overflow: hidden;">
        <div style="width: 100%; height: 100%; background-color: #4B0082; animation: loading 1.5s linear infinite;"></div>
    </div>
</div>
<style>
@keyframes loading {{
    0% {{ transform: translateX(-100%); }}
    100% {{ transform: translateX(100%); }}
}}
</style>
    '''

# -----------------------
# Video Processing Helper
# -----------------------
def downsample_video(video_path):
    """
    Downsamples the video to 10 evenly spaced frames.
    Each frame is converted to a PIL Image along with its timestamp.
    """
    vidcap = cv2.VideoCapture(video_path)
    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    frames = []
    if total_frames <= 0 or fps <= 0:
        vidcap.release()
        return frames
    # Sample 10 evenly spaced frames.
    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
    for i in frame_indices:
        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
        success, image = vidcap.read()
        if success:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(image)
            timestamp = round(i / fps, 2)
            frames.append((pil_image, timestamp))
    vidcap.release()
    return frames

# -----------------------
# Qwen2.5-VL Model (Multimodal)
# -----------------------
MODEL_ID_VL = "Qwen/Qwen2.5-VL-7B-Instruct"  # Alternatively: "Qwen/Qwen2.5-VL-3B-Instruct"
processor = AutoProcessor.from_pretrained(MODEL_ID_VL, trust_remote_code=True)
vl_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_VL,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
).to("cuda").eval()

# -----------------------
# Text Generation Setup (DeepHermes)
# -----------------------
TG_MODEL_ID = "prithivMLmods/DeepHermes-3-Llama-3-3B-Preview-abliterated"
tg_tokenizer = AutoTokenizer.from_pretrained(TG_MODEL_ID)
tg_model = AutoModelForCausalLM.from_pretrained(
    TG_MODEL_ID,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
tg_model.eval()

# -----------------------
# Main Inference Function
# -----------------------
@spaces.GPU
def model_inference(input_dict, history):
    text = input_dict["text"]
    files = input_dict["files"]

    # Video inference branch
    if text.strip().lower().startswith("@video-infer"):
        text = text[len("@video-infer"):].strip()
        if not files:
            yield gr.Error("Please upload a video file along with your @video-infer query.")
            return
        video_path = files[0]
        frames = downsample_video(video_path)
        if not frames:
            yield gr.Error("Could not process video.")
            return
        # Build messages starting with the text prompt and then add each frame with its timestamp.
        messages = [
            {
                "role": "user",
                "content": [{"type": "text", "text": text}]
            }
        ]
        for image, timestamp in frames:
            messages[0]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
            messages[0]["content"].append({"type": "image", "image": image})
        # Collect images from the frames.
        video_images = [image for image, _ in frames]
        prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(
            text=[prompt],
            images=video_images,
            return_tensors="pt",
            padding=True,
        ).to("cuda")
        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
        thread = Thread(target=vl_model.generate, kwargs=generation_kwargs)
        thread.start()
        buffer = ""
        yield progress_bar_html("Processing video with Qwen2.5VL Model")
        for new_text in streamer:
            buffer += new_text
            time.sleep(0.01)
            yield buffer
        return

    # Multimodal branch if images are provided (non-video)
    if files:
        # If more than one file is provided, load them as images.
        if len(files) > 1:
            images = [load_image(image) for image in files]
        elif len(files) == 1:
            images = [load_image(files[0])]
        else:
            images = []

        if text == "":
            yield gr.Error("Please input a text query along with the image(s).")
            return

        messages = [
            {
                "role": "user",
                "content": [
                    *[{"type": "image", "image": image} for image in images],
                    {"type": "text", "text": text},
                ],
            }
        ]
        prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(
            text=[prompt],
            images=images,
            return_tensors="pt",
            padding=True,
        ).to("cuda")
        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
        thread = Thread(target=vl_model.generate, kwargs=generation_kwargs)
        thread.start()
        buffer = ""
        yield progress_bar_html("Processing with Qwen2.5VL Model")
        for new_text in streamer:
            buffer += new_text
            time.sleep(0.01)
            yield buffer
        return

    # Text-only branch using DeepHermes text generation.
    if text.strip() == "":
        yield gr.Error("Please input a query.")
        return

    input_ids = tg_tokenizer(text, return_tensors="pt").to(tg_model.device)
    streamer = TextIteratorStreamer(tg_tokenizer, skip_prompt=True, skip_special_tokens=True)
    generation_kwargs = {
        "input_ids": input_ids,
        "streamer": streamer,
        "max_new_tokens": 2048,
        "do_sample": True,
        "top_p": 0.9,
        "top_k": 50,
        "temperature": 0.6,
        "repetition_penalty": 1.2,
    }
    thread = Thread(target=tg_model.generate, kwargs=generation_kwargs)
    thread.start()
    buffer = ""
    yield progress_bar_html("Processing text with DeepHermes Model")
    for new_text in streamer:
        buffer += new_text
        time.sleep(0.01)
        yield buffer

# -----------------------
# Gradio Chat Interface
# -----------------------
examples = [
    [{"text": "Describe the Image?", "files": ["example_images/document.jpg"]}],
    [{"text": "Tell me a story about a brave knight."}],
    [{"text": "@video-infer Explain the content of the Advertisement", "files": ["example_images/videoplayback.mp4"]}],
    [{"text": "@video-infer Explain the content of the video in detail", "files": ["example_images/breakfast.mp4"]}],
]

demo = gr.ChatInterface(
    fn=model_inference,
    description="# **Qwen2.5-VL-7B-Instruct `@video-infer for video understanding`**",
    examples=examples,
    fill_height=True,
    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image", "video"], file_count="multiple"),
    stop_btn="Stop Generation",
    multimodal=True,
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch(debug=True)