import logging
import mimetypes
from pathlib import Path

import gradio as gr
import spaces
import torch
from PIL import Image
from huggingface_hub import snapshot_download
from transformers import AutoModelForImageTextToText, AutoProcessor

# --- Logging (HF-friendly) ---
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

MODEL_ID = "allenai/Molmo2-8B"

# --- Optional: pre-cache model on CPU startup to reduce first-GPU latency ---
try:
    logger.info(f"Pre-caching {MODEL_ID} (this may take a bit the first time)...")
    snapshot_download(repo_id=MODEL_ID)
    logger.info("Pre-cache complete.")
except Exception as e:
    logger.warning(f"Pre-cache failed (ok): {e}")

# --- Globals for lazy load on GPU ---
processor = None
model = None


def _guess_kind(file_path: str) -> str:
    """
    Returns "video" or "image" based on mime type / extension.
    Defaults to "image" if unsure.
    """
    mime, _ = mimetypes.guess_type(file_path)
    if mime and mime.startswith("video/"):
        return "video"
    if mime and mime.startswith("image/"):
        return "image"

    ext = Path(file_path).suffix.lower()
    if ext in {".mp4", ".mov", ".webm", ".mkv", ".avi"}:
        return "video"
    if ext in {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff"}:
        return "image"

    return "image"


def _build_inputs(file_path: str, prompt: str):
    """
    Builds Molmo2 chat-template inputs for either image or video.
    """
    kind = _guess_kind(file_path)

    if kind == "video":
        content = [
            {"type": "text", "text": prompt},
            {"type": "video", "video": file_path},
        ]
        label = "video"
    else:
        img = Image.open(file_path).convert("RGB")
        content = [
            {"type": "text", "text": prompt},
            {"type": "image", "image": img},
        ]
        label = "image"

    messages = [{"role": "user", "content": content}]

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    return inputs, label


@spaces.GPU(duration=240)  # video can be slower; adjust down if you want tighter limits
def analyze(file_path: str, prompt: str) -> str:
    """
    Unified handler: accepts a filepath (image OR video) + text prompt.
    """
    global processor, model

    if not file_path:
        return "No file uploaded."

    if model is None:
        logger.info("Loading processor + model in ZeroGPU environment...")
        processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
        model = AutoModelForImageTextToText.from_pretrained(
            MODEL_ID,
            trust_remote_code=True,
            torch_dtype="auto",
            device_map="auto",
        )
        model.eval()
        logger.info("Model loaded.")

    inputs, kind = _build_inputs(file_path, prompt)
    logger.info(f"Built inputs for: {kind}")

    # device_map="auto" can shard, but for input tensors we can move to the first param device
    device = next(model.parameters()).device
    inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}

    # safer mixed precision: bf16 if supported, else fp16
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    amp_dtype = torch.bfloat16 if use_bf16 else torch.float16

    with torch.no_grad():
        with torch.autocast("cuda", dtype=amp_dtype):
            generated_ids = model.generate(
                **inputs,
                max_new_tokens=512,   # bump if you want longer
                do_sample=False,
            )

    # decode only the continuation
    prompt_len = inputs["input_ids"].shape[1]
    generated_tokens = generated_ids[0, prompt_len:]
    text_out = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)

    return text_out.strip() or "(No text output.)"


# --- Gradio UI ---
demo = gr.Interface(
    fn=analyze,
    inputs=[
        gr.File(label="Upload image or video", file_types=["image", "video"], type="filepath"),
        gr.Textbox(label="Prompt", placeholder="E.g., 'Describe what happens in this clip' or 'Find the login button'"),
    ],
    outputs=gr.Textbox(label="Model output"),
    title="Molmo 2 8B — Image/Video QA Agent (MCP)",
    description="Upload an image or a short video clip.",
)

if __name__ == "__main__":
    logger.info("Launching Gradio server...")
    demo.launch(server_name="0.0.0.0", server_port=7860, mcp_server=True, show_error=True)