import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch
import torch
import gradio as gr
from threading import Thread
from transformers import AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer
from qwen_vl_utils.vision_process import fetch_video

MODELS = {
    "Code-as-World-VL-9B (recommended)": "MirroS-Lab/Code-as-World-VL-9B",
    "Code-as-World-VL-4B (faster)": "MirroS-Lab/Code-as-World-VL-4B",
}

SYSTEM_PROMPT = (
    "You are an expert video analyst specializing in physics measurements. "
    "Analyze the video frames carefully and provide ONLY the numerical answer with units. "
    "No explanation or reasoning needed. Format your response as: [value] [unit]. "
    "Example: 2.5 cm. Be as accurate as possible with measurements and calculations. "
    "Please give me an estimated answer even if you are not sure."
)

VIDEO_NFRAMES = 16
MAX_NEW_TOKENS = 512
DEFAULT_MODEL_LABEL = list(MODELS)[0]
DEPTH_PREFIX = (
    "Additionally, you have the following information about the distance between "
    "the objects in the video and the shooting camera:"
)

_CACHE = {}


def _prior_with_depth(prior_info, depth_info):
    prior_info = prior_info.strip().rstrip(".")
    depth_info = depth_info.strip().rstrip(".")
    return f"{prior_info}. {DEPTH_PREFIX} {depth_info}"


EXAMPLES = [
    [
        "internet_0005.mp4",
        "What is the height of the woman in yellow in meters?",
        "walking velocity = 1.25 m/s",
    ],
    [
        "captured_0013.mp4",
        "What is the length of the pen in cm?",
        _prior_with_depth(
            "t=1.5, ball acceleration = 3.0 m/s^2",
            (
                "t=0s, distance_cup_camera = 1.2100 m\n"
                "t=0s, distance_pen_camera = 1.1750 m\n"
                "t=0s, distance_cookie_camera = 1.1820 m\n"
                "t=0s, distance_note_book_camera = 0.9620 m\n"
                "t=0s, distance_desk_left_camera = 0.8690 m\n"
                "t=0s, distance_desk_right_camera = 0.8520 m\n"
                "t=0s, distance_slope_far_camera = 1.3450 m\n"
                "t=0s, distance_slope_near_camera = 1.1040 m"
            ),
        ),
    ],
    [
        "simulation_0300.mp4",
        "What is the speed of the car at 1s, in m/s?",
        _prior_with_depth(
            "acceleration of the car = 5 m/s^2",
            (
                "t=1.0s, distance_car_camera = 17.8278 m\n"
                "t=2.0s, distance_car_camera = 23.1989 m"
            ),
        ),
    ],
    [
        "simulation_0193.mp4",
        "What is the diameter of the model of Callisto in meters?",
        "Callisto's model speed = 0.627 m/s",
    ],
]


def _load_model(repo_id):
    if _CACHE.get("id") == repo_id:
        return _CACHE["model"], _CACHE["processor"]
    _CACHE.clear()
    import gc
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    processor = AutoProcessor.from_pretrained(repo_id, trust_remote_code=True)
    model = AutoModelForImageTextToText.from_pretrained(
        repo_id,
        dtype=torch.bfloat16,
        trust_remote_code=True,
    ).to("cuda").eval()
    _CACHE.update(id=repo_id, model=model, processor=processor)
    return model, processor


def _process_video(video_path):
    """Extract frames from a video file using qwen_vl_utils, matching the eval pipeline."""
    vision_info = {
        "video": video_path,
        "min_pixels": 0,
        "max_pixels": 262144,
        "video_fps": 2.0,
        "nframes": VIDEO_NFRAMES,
    }
    try:
        video_out, sample_fps = fetch_video(
            vision_info,
            return_video_sample_fps=True,
            return_video_metadata=True,
        )
    except ValueError as exc:
        import re
        match = re.search(
            r"nframes should in interval \[(\d+), (\d+)\], but got (\d+)", str(exc)
        )
        if match is None:
            raise
        min_allowed, max_allowed, requested = (int(v) for v in match.groups())
        if requested <= max_allowed or max_allowed < min_allowed:
            raise
        vision_info["nframes"] = max_allowed
        video_out, sample_fps = fetch_video(
            vision_info,
            return_video_sample_fps=True,
            return_video_metadata=True,
        )

    if isinstance(video_out, tuple) and len(video_out) == 2:
        video_input, raw_metadata = video_out
    else:
        video_input = video_out
        raw_metadata = None

    num_frames = int(video_input.shape[0]) if hasattr(video_input, "shape") else len(video_input)
    metadata = {}
    if raw_metadata is not None:
        valid_keys = {"total_num_frames", "fps", "width", "height", "duration", "video_backend", "frames_indices"}
        metadata = {k: v for k, v in dict(raw_metadata).items() if k in valid_keys}

    frame_indices = metadata.get("frames_indices")
    if hasattr(frame_indices, "detach"):
        frame_indices = frame_indices.detach().cpu().tolist()
    elif frame_indices is not None:
        frame_indices = list(frame_indices)
    if not frame_indices or len(frame_indices) != num_frames:
        frame_indices = list(range(num_frames))
    raw_fps = float(metadata.get("fps", 0.0) or 0.0)
    metadata["fps"] = raw_fps if raw_fps > 0 else float(
        sample_fps if sample_fps > 0 else 24.0
    )
    metadata["frames_indices"] = frame_indices
    metadata["total_num_frames"] = int(metadata.get("total_num_frames", num_frames))

    return video_input, metadata


def _duration(video_path=None, question="", prior_info="", model_label=DEFAULT_MODEL_LABEL,
              max_new_tokens=MAX_NEW_TOKENS, *args, **kwargs):
    # Measured: cold start ~31s (model load + inference), warm ~9s (inference only)
    load = 0 if _CACHE.get("id") == MODELS.get(model_label, MODELS[DEFAULT_MODEL_LABEL]) else 60
    return int(min(15 + load + 20, 180))


@spaces.GPU(duration=_duration)
def analyze(video=None, question="", prior_info="", model_label=DEFAULT_MODEL_LABEL,
            max_new_tokens=MAX_NEW_TOKENS, progress=gr.Progress(track_tqdm=True)):
    """Analyze a video and answer a quantitative physical reasoning question.

    Args:
        video: Input video file path.
        question: A physics measurement question about the video.
        prior_info: Optional prior information (e.g. "ruler calibre = 1 cm").
        model_label: Which model checkpoint to use.
        max_new_tokens: Maximum number of tokens to generate.
    """
    if not video:
        raise gr.Error("Please upload a video file.")
    if not question:
        raise gr.Error("Please enter a question.")

    question = question.strip()
    prior_info = (prior_info or "").strip()

    # Prior Information may also contain the standard depth-information sentence.
    prefix = ""
    if prior_info:
        prefix = f"Given that {prior_info}"
        if prefix[-1] not in ".!?":
            prefix += "."
    content_text = f"{prefix} {question}".strip() if prefix else question

    model, processor = _load_model(MODELS.get(model_label) or MODELS[DEFAULT_MODEL_LABEL])

    # Process video
    video_input, metadata = _process_video(video)

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": [
            {"type": "video", "video": video_input, "metadata": metadata},
            {"type": "text", "text": content_text + "\n\nPlease answer the question with numbers and units ONLY. No explanation needed."},
        ]},
    ]

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

    streamer = TextIteratorStreamer(
        processor.tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    generation_kwargs = dict(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=False,
        streamer=streamer,
    )
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    output = ""
    for token in streamer:
        output += token
        # Strip thinking tags if present
        if "</think>" in output:
            clean = output.split("</think>", 1)[1].lstrip()
        else:
            clean = output
        yield clean

    thread.join()


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Code-as-World VL Demo") as demo:
    with gr.Row(elem_id="col-container"):
        gr.Markdown(
            "# 🌍 Code-as-World VL: Physical Reasoning from Video\n"
            "Upload a video and ask a quantitative physical reasoning question. "
            "The model estimates measurements like object sizes, velocities, and distances from visual evidence.\n\n"
            "[Paper](https://arxiv.org/abs/2608.27549) | "
            "[GitHub](https://github.com/mirros-lab/code-as-world) | "
            "[9B Model](https://huggingface.co/MirroS-Lab/Code-as-World-VL-9B) | "
            "[4B Model](https://huggingface.co/MirroS-Lab/Code-as-World-VL-4B)"
        )

    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=1):
            video = gr.Video(label="Input Video", sources=["upload"])
            question = gr.Textbox(
                label="Question",
                placeholder="e.g. What is the length of the wood block in cm?",
                lines=2,
            )
            prior_info = gr.Textbox(
                label="Prior Information (optional)",
                placeholder="e.g. ruler calibre = 1 cm; depth context may be included here",
                lines=3,
            )
            with gr.Accordion("Advanced settings", open=False):
                model_label = gr.Dropdown(
                    list(MODELS), value=DEFAULT_MODEL_LABEL, label="Model"
                )
                max_new_tokens = gr.Slider(
                    64, 1024, value=MAX_NEW_TOKENS, step=64,
                    label="Max new tokens",
                )
            run = gr.Button("Analyze", variant="primary")

        with gr.Column(scale=2):
            output = gr.Textbox(
                label="Answer",
                lines=10,
                placeholder="The model's answer will appear here…",
            )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[video, question, prior_info],
    )

    run.click(
        analyze,
        inputs=[video, question, prior_info, model_label, max_new_tokens],
        outputs=output,
        api_name="analyze",
    )

demo.queue(max_size=10).launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)