#!/usr/bin/env python

import os
import re
import tempfile
from collections.abc import Iterator
from threading import Thread
import cv2
import gradio as gr
import spaces
import torch
from loguru import logger
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer

model_id = os.getenv("MODEL_ID", "google/medgemma-4b-it")
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForImageTextToText.from_pretrained(
    model_id, device_map="auto", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,token=os.environ.get("HF_TOKEN", "YOUR_HF_TOKEN")
)

MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))

def count_files_in_new_message(paths: list[str]) -> tuple[int, int]:
    image_count = 0
    video_count = 0
    for path in paths:
        if path.endswith(".mp4"):
            video_count += 1
        else:
            image_count += 1
    return image_count, video_count

def count_files_in_history(history: list[dict]) -> tuple[int, int]:
    image_count = 0
    video_count = 0
    for item in history:
        if item["role"] != "user" or isinstance(item["content"], str):
            continue
        if item["content"][0].endswith(".mp4"):
            video_count += 1
        else:
            image_count += 1
    return image_count, video_count

def validate_media_constraints(message: dict, history: list[dict]) -> bool:
    new_image_count, new_video_count = count_files_in_new_message(message["files"])
    history_image_count, history_video_count = count_files_in_history(history)
    image_count = history_image_count + new_image_count
    video_count = history_video_count + new_video_count
    if video_count > 1:
        gr.Warning("Only one video is supported.")
        return False
    if video_count == 1:
        if image_count > 0:
            gr.Warning("Mixing images and videos is not allowed.")
            return False
        if "<image>" in message["text"]:
            gr.Warning("Using <image> tags with video files is not supported.")
            return False
    if video_count == 0 and image_count > MAX_NUM_IMAGES:
        gr.Warning(f"You can upload up to {MAX_NUM_IMAGES} images.")
        return False
    if "<image>" in message["text"] and message["text"].count("<image>") != new_image_count:
        gr.Warning("The number of <image> tags in the text does not match the number of images.")
        return False
    return True

def downsample_video(video_path: str) -> list[tuple[Image.Image, float]]:
    vidcap = cv2.VideoCapture(video_path)
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))

    frame_interval = max(total_frames // MAX_NUM_IMAGES, 1)
    frames: list[tuple[Image.Image, float]] = []

    for i in range(0, min(total_frames, MAX_NUM_IMAGES * frame_interval), frame_interval):
        if len(frames) >= MAX_NUM_IMAGES:
            break

        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

def process_video(video_path: str) -> list[dict]:
    content = []
    frames = downsample_video(video_path)
    for frame in frames:
        pil_image, timestamp = frame
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
            pil_image.save(temp_file.name)
            content.append({"type": "text", "text": f"Frame {timestamp}:"})
            content.append({"type": "image", "url": temp_file.name})
    logger.debug(f"{content=}")
    return content

def process_interleaved_images(message: dict) -> list[dict]:
    logger.debug(f"{message['files']=}")
    parts = re.split(r"(<image>)", message["text"])
    logger.debug(f"{parts=}")

    content = []
    image_index = 0
    for part in parts:
        logger.debug(f"{part=}")
        if part == "<image>":
            content.append({"type": "image", "url": message["files"][image_index]})
            logger.debug(f"file: {message['files'][image_index]}")
            image_index += 1
        elif part.strip():
            content.append({"type": "text", "text": part.strip()})
        elif isinstance(part, str) and part != "<image>":
            content.append({"type": "text", "text": part})
    logger.debug(f"{content=}")
    return content

def process_new_user_message(message: dict) -> list[dict]:
    if not message["files"]:
        return [{"type": "text", "text": message["text"]}]

    if message["files"][0].endswith(".mp4"):
        return [{"type": "text", "text": message["text"]}, *process_video(message["files"][0])]

    if "<image>" in message["text"]:
        return process_interleaved_images(message)

    return [
        {"type": "text", "text": message["text"]},
        *[{"type": "image", "url": path} for path in message["files"]],
    ]

def process_history(history: list[dict]) -> list[dict]:
    messages = []
    current_user_content: list[dict] = []
    for item in history:
        if item["role"] == "assistant":
            if current_user_content:
                messages.append({"role": "user", "content": current_user_content})
                current_user_content = []
            messages.append({"role": "assistant", "content": [{"type": "text", "text": item["content"]}]})
        else:
            content = item["content"]
            if isinstance(content, str):
                current_user_content.append({"type": "text", "text": content})
            else:
                current_user_content.append({"type": "image", "url": content[0]})
    return messages

@spaces.GPU(duration=120)
def run(message: dict, history: list[dict], system_prompt: str = "", max_new_tokens: int = 2048) -> Iterator[str]:
    if not validate_media_constraints(message, history):
        yield ""
        return

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})
    messages.extend(process_history(history))
    messages.append({"role": "user", "content": process_new_user_message(message)})

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(device=model.device)

    streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        inputs,
        max_new_tokens=max_new_tokens,
        streamer=streamer,
        temperature=1.0,
        top_p=0.95,
        top_k=64,
        min_p=0.0,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    output = ""
    for delta in streamer:
        output += delta
        yield output

# Custom CSS for the UI
custom_css = """
:root {
    --primary: #4f46e5;
    --primary-dark: #4338ca;
    --text: #1f2937;
    --background: #f9fafb;
    --chat-bg: #ffffff;
    --user-bubble: #e0e7ff;
    --bot-bubble: #f3f4f6;
    --border: #e5e7eb;
}

.dark {
    --primary: #6366f1;
    --primary-dark: #4f46e5;
    --text: #f3f4f6;
    --background: #111827;
    --chat-bg: #1f2937;
    --user-bubble: #4338ca;
    --bot-bubble: #374151;
    --border: #4b5563;
}

body {
    font-family: 'Inter', sans-serif;
}

.gr-chatbot {
    background-color: var(--chat-bg);
    border-radius: 12px;
    border: 1px solid var(--border);
    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
}

.gr-chat-message {
    padding: 16px 20px;
    border-radius: 12px;
    margin: 8px 0;
    max-width: 80%;
}

.gr-chat-message-user {
    background-color: var(--user-bubble);
    margin-left: auto;
    border-bottom-right-radius: 4px;
}

.gr-chat-message-bot {
    background-color: var(--bot-bubble);
    margin-right: auto;
    border-bottom-left-radius: 4px;
}

.gr-textbox textarea {
    min-height: 120px;
    border-radius: 12px;
    padding: 16px;
    background-color: var(--background);
    color: var(--text);
    border: 1px solid var(--border);
}

.gr-button {
    background-color: var(--primary) !important;
    color: white !important;
    border-radius: 8px !important;
    padding: 10px 20px !important;
    font-weight: 500 !important;
    transition: all 0.2s !important;
}

.gr-button:hover {
    background-color: var(--primary-dark) !important;
    transform: translateY(-1px) !important;
}

.gr-button:active {
    transform: translateY(0) !important;
}

.gr-interface {
    max-width: 900px;
    margin: 0 auto;
    padding: 24px;
}

.gr-header {
    text-align: center;
    margin-bottom: 24px;
}

.gr-header h1 {
    font-size: 2.5rem;
    font-weight: 700;
    color: var(--primary);
    margin-bottom: 8px;
}

.gr-header p {
    color: var(--text);
    opacity: 0.8;
    font-size: 1.1rem;
}

.gr-image-preview {
    border-radius: 8px;
    max-width: 100%;
    max-height: 300px;
    object-fit: contain;
}

.gr-video-preview {
    border-radius: 8px;
    max-width: 100%;
    max-height: 300px;
}
"""

DESCRIPTION = """\
## Medical Vision-Language Assistant

This advanced AI assistant can understand and analyze medical images, videos, and text. 
Upload images or a video along with your questions to get insights.

**Features:**
- Analyze medical images (X-rays, CT scans, etc.)
- Process video frames from medical videos
- Interleave images with text questions
- Customize system behavior with prompts
"""

demo = gr.ChatInterface(
    fn=run,
    type="messages",
    chatbot=gr.Chatbot(
        type="messages", 
        scale=1, 
        allow_tags=["image"],
        bubble_full_width=False,
        avatar_images=(
            "assets/user.png",  # User avatar
            "assets/doctor.png"  # Bot avatar (replace with your own)
        ),
        render=False  # We'll handle rendering in CSS
    ),
    textbox=gr.MultimodalTextbox(
        file_types=["image", ".mp4"], 
        file_count="multiple", 
        autofocus=True,
        placeholder="Type your message or upload images/video...",
    ),
    multimodal=True,
    additional_inputs=[
        gr.Textbox(
            label="System Prompt", 
            value="You are a helpful and knowledgeable medical expert. Provide accurate, detailed explanations in clear language.",
            info="Guide the assistant's behavior and expertise"
        ),
        gr.Slider(
            label="Response Length", 
            minimum=100, 
            maximum=4096, 
            step=10, 
            value=1024,
            info="Control how verbose the responses are"
        ),
    ],
    stop_btn=None,
    title="",
    description=DESCRIPTION,
    examples=[
        ["What abnormalities do you see in this chest X-ray?", "examples/chest_xray.jpg"],
        ["Explain the key findings in this MRI scan.", "examples/brain_mri.jpg"],
        ["Describe the progression shown in this video.", "examples/heart_ultrasound.mp4"],
    ],
    cache_examples=False,
    css=custom_css,
    theme=gr.themes.Default(
        primary_hue="indigo",
        secondary_hue="gray",
        font=["Inter", "sans-serif"]
    ),
)

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