import os

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

import spaces  # noqa: E402  (must come before torch)

import base64  # noqa: E402
import difflib  # noqa: E402
import hashlib  # noqa: E402
import io  # noqa: E402
import json  # noqa: E402
import re  # noqa: E402
import shutil  # noqa: E402
import signal  # noqa: E402
import subprocess  # noqa: E402
import sys  # noqa: E402
import tempfile  # noqa: E402
import threading  # noqa: E402
import time  # noqa: E402
import uuid  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from jinja2 import StrictUndefined, Template  # noqa: E402
from PIL import Image  # noqa: E402
from transformers import AutoModelForImageTextToText, AutoProcessor, TextIteratorStreamer  # noqa: E402

ROOT = Path(__file__).resolve().parent
HARNESS_ROOT = ROOT / "harness"
sys.path.insert(0, str(HARNESS_ROOT))
for _tool in (HARNESS_ROOT / "bin").iterdir():  # the Hub does not keep the executable bit
    _tool.chmod(0o755)

from video_coder_harness.blender_exec import execute as blender_execute  # noqa: E402
from video_coder_harness.guard import (  # noqa: E402
    BLOCKED,
    REPEAT_WARNING,
    _normalize_output,
    harness_violation,
)
from video_coder_harness.static_check import check_file  # noqa: E402

MODEL_ID = "Video-Coder/Video-Coder-9B"

# --------------------------------------------------------------------------------------------------------------------
# Blender 4.5 (bpy 4.5.14 on Python 3.11, same as the authors' bpy-agent.sif sandbox), rendered on CPU with Mesa
# llvmpipe. Installed into its own venv at startup because the Space itself runs Python 3.12.
# --------------------------------------------------------------------------------------------------------------------
BPY_ENV = Path(os.environ.get("BPY_ENV", str(Path.home() / "bpyenv")))
BPY_PY = BPY_ENV / "bin" / "python"


def setup_bpy() -> None:
    if BPY_PY.exists():
        r = subprocess.run([str(BPY_PY), "-c", "import bpy; print(bpy.app.version_string)"],
                           capture_output=True, text=True)
        if r.returncode == 0:
            print("bpy ready:", r.stdout.strip())
            return
    t0 = time.time()
    subprocess.run([sys.executable, "-m", "uv", "venv", "-q", "-p", "3.11", "--clear", str(BPY_ENV)], check=True)
    subprocess.run([sys.executable, "-m", "uv", "pip", "install", "-q", "-p", str(BPY_PY), "bpy==4.5.14"],
                   check=True)
    r = subprocess.run([str(BPY_PY), "-c", "import bpy; print(bpy.app.version_string)"],
                       capture_output=True, text=True)
    print(f"bpy install took {time.time() - t0:.0f}s:", r.stdout.strip(), r.stderr[-2000:])


setup_bpy()

# Harness settings: configs/teacher.yaml + configs/uni_container_visual.yaml (the RL / inference setting).
HARNESS = {
    "blender_version": "4.5",
    "blender_available": True,
    "bpy_python": str(BPY_PY),
    "bpy_libs": "/usr",
    "exec_timeout": 600,
    "summary_samples": 6,
    "content_samples": 8,
    "visual_feedback": True,
    "preview_frames": 8,
    "preview_width": 640,
    "judge": {"enabled": False},
    "output_video": "output.mp4",
    "fps": 24,
    "frame_count": 120,
    "max_frame_count": 240,
    "resolution": [1280, 720],
    "script_name": "main.py",
}
SCRIPT = HARNESS["script_name"]
CMD_TIMEOUT = 900
REPEAT_WARN, REPEAT_STOP, REPEAT_SIM = 3, 6, 0.9
MAX_FORMAT_ERRORS = 3
RUNS_DIR = Path(tempfile.gettempdir()) / "video_coder_runs"
RUNS_DIR.mkdir(parents=True, exist_ok=True)
_SECRET_ENV = re.compile(r"(KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL)", re.I)
HARNESS_MARKERS = [str(HARNESS_ROOT), str(BPY_ENV)]

SYSTEM_PROMPT = (
    "You are an expert Blender technical artist and Python programmer. You write Blender Python (bpy) scripts that "
    "build and animate 3D scenes and render them as untextured white-model (clay) videos. You work in a Linux shell "
    "through a bash tool."
)
# The exact user prompt of the SFT trajectories (Video-Coder/Train-Data sft_data.json), with the scene as __TASK__.
INSTANCE_TEMPLATE = (HARNESS_ROOT / "instance_template.txt").read_text()
TOOLS = [{
    "type": "function",
    "function": {
        "name": "bash",
        "description": "Execute a bash command",
        "parameters": {
            "type": "object",
            "properties": {"command": {"type": "string", "description": "The bash command to execute"}},
            "required": ["command"],
        },
    },
}]
OBSERVATION_TEMPLATE = Template("""{%- if output.output | length < 10000 -%}
{
  "returncode": {{ output.returncode }},
  "output": {{ output.output | tojson }}
  {%- if output.exception_info %}, "exception_info": {{ output.exception_info | tojson }}{% endif %}
}
{%- else -%}
{
  "returncode": {{ output.returncode }},
  "output_head": {{ output.output[:5000] | tojson }},
  "output_tail": {{ output.output[-5000:] | tojson }},
  "elided_chars": {{ output.output | length - 10000 }},
  "warning": "Output too long."
  {%- if output.exception_info %}, "exception_info": {{ output.exception_info | tojson }}{% endif %}
}
{%- endif -%}
""", undefined=StrictUndefined)
FORMAT_ERROR_TEMPLATE = Template("""{% if finish_reason is defined and (finish_reason == "length" or (finish_reason == "tool_calls" and not has_tool_calls)) -%}
Your previous response reached the output token limit (finish_reason={{ finish_reason }}) before you produced a tool call, so it was cut off. Respond more concisely and finish with exactly one bash tool call.
{%- else -%}
Tool call error:

<error>
{{error}}
</error>

Every response must call the 'bash' tool exactly once with {"command": "..."}.
To finish, run `echo COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT` on its own.
{%- endif %}
""")
MM_TAG = re.compile(r"(?s)<MSWEA_MULTIMODAL_CONTENT><CONTENT_TYPE>(.+?)</CONTENT_TYPE>(.+?)</MSWEA_MULTIMODAL_CONTENT>")
TOOL_CALL_RE = re.compile(r"<tool_call>\s*<function=([^>\s]+)>(.*?)</function>\s*</tool_call>", re.S)
PARAM_RE = re.compile(r"<parameter=([^>\s]+)>\n?(.*?)\n?</parameter>", re.S)
END_TOKENS = ("<|im_end|>", "<|endoftext|>")

# --------------------------------------------------------------------------------------------------------------------
# Model: Video-Coder-9B (Qwen3.5-9B, vision-language) in bf16, fully on the GPU.
# --------------------------------------------------------------------------------------------------------------------
processor = AutoProcessor.from_pretrained(MODEL_ID)
tokenizer = processor.tokenizer
model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to("cuda").eval()

GPU_CHUNK_SECONDS = 90  # one GPU call generates for at most this long; long turns continue in the next call


@spaces.GPU(duration=120)
def llm_chunk(messages: list, images: list, prefix: str = "", max_new_tokens: int = 16384,
              temperature: float = 1.0, top_p: float = 0.95):
    """Continue the agent's current turn for up to GPU_CHUNK_SECONDS; yields (text so far, finished, n_tokens)."""
    prompt = tokenizer.apply_chat_template(messages, tools=TOOLS, add_generation_prompt=True, tokenize=False)
    prompt += prefix
    inputs = processor(text=[prompt], images=images or None, return_tensors="pt").to("cuda")
    n_in = inputs["input_ids"].shape[1]
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
    result = {}

    def _run():
        try:
            out = model.generate(**inputs, streamer=streamer, max_new_tokens=int(max_new_tokens), do_sample=True,
                                 temperature=float(temperature), top_p=float(top_p), top_k=0, use_cache=True,
                                 max_time=GPU_CHUNK_SECONDS)
            result["n"] = int(out.shape[1] - n_in)
        except Exception as e:  # surfaced below
            result["error"] = f"{type(e).__name__}: {e}"
            streamer.end()

    th = threading.Thread(target=_run, daemon=True)
    th.start()
    text = ""
    last = 0.0
    for piece in streamer:
        text += piece
        if time.time() - last > 0.4:
            last = time.time()
            yield text, False, None
    th.join()
    if "error" in result:
        raise gr.Error(result["error"])
    finished = any(t in text for t in END_TOKENS)
    yield text, finished, result.get("n", 0)


# --------------------------------------------------------------------------------------------------------------------
# Agent loop: mini-swe-agent DefaultAgent + video_coder_harness GuardedAgent semantics, re-implemented without the
# litellm/vLLM gateway (the model is called in-process).
# --------------------------------------------------------------------------------------------------------------------
def agent_env() -> dict:
    env = dict(os.environ)
    env.update({k: "" for k in os.environ if _SECRET_ENV.search(k)})
    env.update({"PAGER": "cat", "MANPAGER": "cat", "PIP_PROGRESS_BAR": "off", "TQDM_DISABLE": "1"})
    env["PATH"] = f"{HARNESS_ROOT / 'bin'}:{os.environ.get('PATH', '')}"
    env["VCSFT_PYTHON"] = sys.executable
    env["VCSFT_HARNESS"] = json.dumps(HARNESS)
    return env


def run_command(command: str, cwd: str) -> dict:
    """mini-swe-agent LocalEnvironment.execute."""
    proc = subprocess.Popen(command, shell=True, text=True, cwd=cwd, env=agent_env(), encoding="utf-8",
                            errors="replace", stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                            start_new_session=True)
    try:
        stdout, _ = proc.communicate(timeout=CMD_TIMEOUT)
        return {"output": stdout, "returncode": proc.returncode, "exception_info": ""}
    except subprocess.TimeoutExpired as e:
        os.killpg(proc.pid, signal.SIGKILL)
        stdout, _ = proc.communicate()
        return {"output": stdout or "", "returncode": -1,
                "exception_info": f"An error occurred while executing the command: {e}"}


def is_submit(output: dict) -> bool:
    lines = output.get("output", "").lstrip().splitlines()
    return bool(lines) and lines[0].strip() == "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" and output["returncode"] == 0


def split_images(text: str):
    """video_coder_harness.models.split_images: cut frame images out of the tool output."""
    images = []

    def repl(m):
        if m.group(1).strip() != "image_url":
            return m.group(0)
        before = text[: m.start()].rstrip().splitlines()
        caption = before[-1].strip() if before else ""
        images.append((caption, m.group(2).strip()))
        return f"[image {len(images)} attached below]"

    return MM_TAG.sub(repl, text), images


def data_url_to_pil(url: str) -> Image.Image:
    return Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))).convert("RGB")


def parse_response(raw: str):
    """Split a Qwen3.5 response into (reasoning, content, tool calls) like vLLM's qwen3 reasoning + qwen3_coder parsers."""
    text = raw
    for t in END_TOKENS:
        text = text.split(t)[0]
    reasoning = ""
    if "</think>" in text:
        reasoning, text = text.split("</think>", 1)
        reasoning = reasoning.replace("<think>", "").strip()
    calls = []
    for m in TOOL_CALL_RE.finditer(text):
        args = {k: v for k, v in PARAM_RE.findall(m.group(2))}
        calls.append({"name": m.group(1).strip(), "arguments": args})
    content = TOOL_CALL_RE.sub("", text)
    content = content.split("<tool_call>")[0].strip()
    return reasoning, content, calls


def sha(path: Path) -> str | None:
    return hashlib.sha256(path.read_bytes()).hexdigest() if path.is_file() else None


def short(text: str, n: int = 60) -> str:
    lines = text.splitlines()
    return text if len(lines) <= n else "\n".join(lines[:n]) + f"\n... ({len(lines) - n} more lines)"


def read_script(workdir: Path) -> str:
    p = workdir / SCRIPT
    return p.read_text(errors="replace") if p.is_file() else ""


def new_state(prompt: str, max_steps: int, temperature: float, top_p: float, max_new_tokens: int) -> dict:
    workdir = (RUNS_DIR / uuid.uuid4().hex).resolve()
    workdir.mkdir(parents=True)
    return {
        "workdir": str(workdir),
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": INSTANCE_TEMPLATE.replace("__TASK__", prompt.strip())},
        ],
        "image_urls": [],  # frames shown to the model, in order (JPEG data URLs)
        "chat": [],
        "video": None,
        "preview": None,
        "last_ok": None,  # (script hash, video) of the last run_blender that reported RUN OK
        "last_key": None,
        "repeats": 0,
        "n_format_errors": 0,
        "step": 1,
        "raw": "",  # current (possibly unfinished) model turn
        "used": 0,
        "max_steps": int(max_steps),
        "temperature": float(temperature),
        "top_p": float(top_p),
        "max_new_tokens": int(max_new_tokens),
        "done": False,
        "status": "Starting…",
        "t_start": time.time(),
    }


def ui(state: dict, tick=None):
    workdir = Path(state["workdir"])
    msg = f"**{state['status']}** · {time.time() - state['t_start']:.0f}s elapsed"
    return (state, state["chat"], state["video"], read_script(workdir), state["preview"], msg,
            gr.update() if tick is None else tick)


def start(prompt: str, max_steps: int = 16, temperature: float = 1.0, top_p: float = 0.95,
          max_new_tokens: int = 16384):
    """Start a Video-Coder agent run for a scene description. The run then proceeds turn by turn: call `advance`
    with the returned state (same session) until its status says it finished.

    Args:
        prompt: description of the scene / video to build.
        max_steps: maximum number of agent turns (one bash command each).
        temperature: sampling temperature (the authors sample with 1.0).
        top_p: nucleus sampling (the authors use 0.95).
        max_new_tokens: token budget of a single agent turn.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a scene description.")
    state = new_state(prompt, max_steps, temperature, top_p, max_new_tokens)
    return ui(state, tick=time.time())


def _finish_turn(state: dict):
    """Parse the finished model turn, run its command(s) in the harness; a generator of UI updates."""
    workdir = Path(state["workdir"])
    step, raw = state["step"], state["raw"]
    chat, messages = state["chat"], state["messages"]
    if chat and chat[-1].get("metadata", {}).get("status") == "pending":
        chat.pop()
    reasoning, content, calls = parse_response(raw)
    state["raw"], state["used"] = "", 0
    if reasoning:
        chat.append({"role": "assistant", "content": reasoning,
                     "metadata": {"title": f"💭 Step {step}: reasoning", "status": "done"}})
    msg = {"role": "assistant", "content": content}
    if reasoning:
        msg["reasoning_content"] = reasoning
    if calls:
        msg["tool_calls"] = [{"type": "function", "function": {"name": c["name"], "arguments": c["arguments"]}}
                             for c in calls]
    messages.append(msg)

    error = None
    if not calls:
        error = "No tool calls found in the response. Every response MUST include at least one tool call."
    elif any(c["name"] != "bash" or "command" not in c["arguments"] for c in calls):
        error = "Unknown tool or missing 'command' argument."
    if error:
        state["n_format_errors"] += 1
        finish_reason = "length" if not any(t in raw for t in END_TOKENS) else "stop"
        err = FORMAT_ERROR_TEMPLATE.render(error=error, finish_reason=finish_reason, has_tool_calls=bool(calls))
        if content:
            chat.append({"role": "assistant", "content": content})
        chat.append({"role": "user", "content": f"⚠️ {err}"})
        messages.append({"role": "user", "content": err})
        if state["n_format_errors"] >= MAX_FORMAT_ERRORS:
            return "FormatError"
        return None
    state["n_format_errors"] = 0

    for call in calls:
        command = call["arguments"]["command"]
        shown = f"{content}\n\n" if content else ""
        chat.append({"role": "assistant", "content": f"{shown}```bash\n{short(command)}\n```"})
        first = command.strip().splitlines()[0][:60] if command.strip() else ""
        state["status"] = f"Step {step}: running `{first}`"
        yield
        reason = harness_violation(command, HARNESS_MARKERS, str(workdir))
        if reason:
            out = {"output": BLOCKED.format(reason=reason), "returncode": 1, "exception_info": ""}
        else:
            out = run_command(command, str(workdir))
            if is_submit(out):
                return "Submitted"
        key = (command, _normalize_output(out.get("output", "")))
        last = state["last_key"]
        same = (last is not None and key[1] == last[1]
                and difflib.SequenceMatcher(None, key[0], last[0]).ratio() >= REPEAT_SIM)
        state["repeats"] = state["repeats"] + 1 if same else 1
        state["last_key"] = key
        if state["repeats"] >= REPEAT_WARN:
            out["output"] = (out.get("output") or "") + REPEAT_WARNING.format(n=state["repeats"], stop=REPEAT_STOP)

        text, imgs = split_images(out.get("output") or "")
        obs = OBSERVATION_TEMPLATE.render(output={**out, "output": text})
        if imgs:
            parts = [{"type": "text", "text": obs}]
            for k, (caption, url) in enumerate(imgs, 1):
                parts += [{"type": "text", "text": f"image {k}: {caption}" if caption else f"image {k}:"},
                          {"type": "image"}]
                state["image_urls"].append(url)
            messages.append({"role": "tool", "content": parts})
        else:
            messages.append({"role": "tool", "content": obs})

        if command.strip().startswith("run_blender"):
            vid = workdir / HARNESS["output_video"]
            if vid.is_file():
                keep = workdir / f"render_step{step}.mp4"
                shutil.copy(vid, keep)
                state["video"] = str(keep)
                if text.startswith("RUN OK"):
                    state["last_ok"] = (sha(workdir / SCRIPT), str(keep))
            if (workdir / "preview.png").is_file():
                keep_png = workdir / f"preview_step{step}.png"
                shutil.copy(workdir / "preview.png", keep_png)
                state["preview"] = str(keep_png)
        chat.append({"role": "assistant", "content": f"```\n{short(text.strip(), 40)}\n```",
                     "metadata": {"title": f"🛠️ Output (exit code {out['returncode']})", "status": "done"}})
        if state["repeats"] >= REPEAT_STOP:
            return "RepeatedActions"
    return None


def _accept(state: dict):
    """Harness acceptance: static check + the submitted script re-rendered from scratch (final_execution)."""
    workdir = Path(state["workdir"])
    script = workdir / SCRIPT
    state["chat"].append({"role": "assistant", "content": "✅ Submitted. Validating the final script…"})
    if not script.is_file():
        state["status"] = f"Submitted, but no {SCRIPT} was written"
        return
    check = check_file(script)
    if not check["ok"]:
        state["status"] = f"Submitted, but the static check failed: {check['errors']}"
        return
    last_ok = state["last_ok"]
    if last_ok and last_ok[0] == sha(script):  # Workbench rendering is deterministic: reuse the identical render
        state["video"] = last_ok[1]
    else:
        state["status"] = "Re-rendering the submitted script"
        yield
        final_dir = workdir / "final_run"
        final_dir.mkdir(exist_ok=True)
        shutil.copy(script, final_dir / SCRIPT)
        report = blender_execute(final_dir / SCRIPT, final_dir, HARNESS)
        if (final_dir / HARNESS["output_video"]).is_file():
            state["video"] = str(final_dir / HARNESS["output_video"])
        if (final_dir / "preview.png").is_file():
            state["preview"] = str(final_dir / "preview.png")
        if not report["ok"]:
            state["status"] = f"Submitted, but the final render has problems: {report['problems']}"
            return
    state["status"] = "Accepted ✅ — draft video rendered from the agent's Blender script"


def advance(state: dict):
    """Run the next piece of the agent: one GPU generation chunk of the current turn and, once the turn is complete,
    its bash command in the Video-Coder harness. Returns a new tick while the run is still going.

    Args:
        state: the run state returned by `start` / the previous `advance`.
    """
    if not state or state.get("done"):
        yield ui(state) if state else (state, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(),
                                       gr.update())
        return
    step = state["step"]
    chat = state["chat"]
    if not (chat and chat[-1].get("metadata", {}).get("status") == "pending"):
        chat.append({"role": "assistant", "content": "",
                     "metadata": {"title": f"💭 Step {step}: thinking…", "status": "pending"}})
    live = chat[-1]
    state["status"] = f"Step {step}/{state['max_steps']}: Video-Coder is thinking"
    yield ui(state)

    # ---- one GPU call (the ZeroGPU token of a request is short-lived, so each request makes exactly one) ----
    images = [data_url_to_pil(u) for u in state["image_urls"]]
    base = state["raw"]
    finished, n = False, None
    t0 = time.time()
    for text, finished, n in llm_chunk(state["messages"], images, base, state["max_new_tokens"] - state["used"],
                                       state["temperature"], state["top_p"]):
        state["raw"] = base + text
        live["content"] = state["raw"][-3000:].replace("<think>", "")
        state["status"] = f"Step {step}/{state['max_steps']}: Video-Coder is writing ({len(state['raw'])} chars)"
        yield ui(state)
    state["used"] += n or 0
    print(f"step {step}: chunk {time.time() - t0:.1f}s, {n} tokens, finished={finished}")
    if not finished and state["used"] < state["max_new_tokens"]:
        yield ui(state, tick=time.time())  # continue this turn in the next request
        return

    # ---- the turn is complete: execute it ----
    exit_status = None
    gen = _finish_turn(state)
    while True:
        try:
            next(gen)
            yield ui(state)
        except StopIteration as stop:
            exit_status = stop.value
            break
    state["step"] += 1
    if exit_status is None and state["step"] > state["max_steps"]:
        exit_status = "StepLimit"
    if exit_status is None:
        state["status"] = f"Step {step}/{state['max_steps']} done"
        yield ui(state, tick=time.time())
        return
    state["done"] = True
    if exit_status == "Submitted":
        for _ in _accept(state):
            yield ui(state)
    else:
        state["status"] = {
            "StepLimit": "Step limit reached before the agent submitted (showing the latest render)",
            "RepeatedActions": "Stopped: the agent repeated the same command",
            "FormatError": "Stopped after repeated format errors",
        }[exit_status]
    yield ui(state)


def stop_run(state: dict):
    """Stop the current run."""
    if state:
        state["done"] = True
        state["status"] = "Stopped"
        return ui(state)
    return (state, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update())


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

EXAMPLES = [
    "The video captures a small bird perched on a tree branch. The bird, with its brown and white feathers, yellow "
    "beak, and black and white head, is the main focus of the video. The bird is seen in three different positions on "
    "the branch, each time facing a different direction. The background of the video is filled with green foliage, "
    "providing a natural and serene setting for the bird. The style of the video is a close-up shot.",
    "A single polished walnut-wood chess pawn stands on a small square pedestal of white Carrara marble. No other "
    "objects are in the scene. The camera makes one complete 360-degree orbit around the pawn while slowly rising.",
    "The video shows a sequence of three identical white cars driving on a road. The first car is in the foreground, "
    "the second car is slightly behind and to the right, and the third car is in the background. The cars are moving "
    "from left to right across the frame. The road appears to be a city street with a dark grey surface. The "
    "background features a city skyline with buildings of various heights, and the sky is a gradient of orange and "
    "yellow, suggesting either sunrise or sunset. The style of the video is realistic with a focus on the cars and "
    "their movement.",
    "a man standing in front of a machine.",
    "The video captures a lively scene on the ocean, with three boats in motion. The first boat, a large white yacht, "
    "is seen in the distance, sailing towards the horizon. The second and third boats are closer to the viewer, both "
    "moving at high speeds, leaving trails of white water behind them. The ocean is a deep blue, and the sky above is "
    "clear, suggesting a sunny day. The boats are the main focus of the video, with their movement and the water's "
    "reaction to their speed creating a dynamic and exciting scene. The overall style of the video is a dynamic "
    "aerial shot, capturing the movement and beauty of the ocean and the boats.",
]

with gr.Blocks(title="Video-Coder") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🎬 Video-Coder: agentic coding for video generation\n"
            "**Video-Coder-9B** writes and runs a Blender Python program for your scene description, inspects the "
            "rendered frames, fixes what is wrong and submits a **draft video** (white-model 3D render with the "
            "layout, motion and camera of the scene) — the structure/motion signal the paper feeds to a downstream "
            "video generator.\n\n"
            "[📖 Paper](https://huggingface.co/papers/2611.09894) · [🌐 Project page](https://cywang735.github.io/"
            "Video-Coder/) · [💻 Code](https://github.com/CYWang735/Video-Coder) · [🤗 Model](https://huggingface.co/"
            "Video-Coder/Video-Coder-9B)\n\n"
            "⏳ A full agent run takes several minutes: the first turn writes a long script, and every `run_blender` "
            "renders the animation with Blender on CPU. The trajectory streams live; keep this tab open."
        )
        with gr.Row():
            prompt = gr.Textbox(label="Scene description", lines=3, scale=4,
                                placeholder="Describe the scene and its motion, e.g. 'a red ball rolls down a ramp "
                                            "and knocks over three dominoes'")
            with gr.Column(scale=1, min_width=180):
                run_btn = gr.Button("Generate draft video", variant="primary")
                stop_btn = gr.Button("Stop", variant="stop")
        status = gr.Markdown()
        with gr.Row():
            with gr.Column(scale=1):
                video_out = gr.Video(label="Draft video (Blender render)", autoplay=True)
                preview_out = gr.Image(label="Latest render: frames 1, ⅓, ⅔, last", type="filepath")
            with gr.Column(scale=1):
                with gr.Tabs():
                    with gr.Tab("Agent trajectory"):
                        chat_out = gr.Chatbot(label="Video-Coder harness", height=640, show_label=False)
                    with gr.Tab("main.py"):
                        code_out = gr.Code(label="Blender script written by the agent", language="python",
                                           lines=30)
        with gr.Accordion("Advanced settings", open=False):
            max_steps = gr.Slider(2, 40, value=16, step=1, label="Max agent steps (one bash command each)")
            temperature = gr.Slider(0.1, 1.5, value=1.0, step=0.05, label="Temperature")
            top_p = gr.Slider(0.5, 1.0, value=0.95, step=0.01, label="Top-p")
            max_new_tokens = gr.Slider(2048, 32768, value=16384, step=1024, label="Max tokens per agent turn")

        state = gr.State(None)
        tick = gr.Number(value=0, visible=False)
        outputs = [state, chat_out, video_out, code_out, preview_out, status, tick]
        inputs = [prompt, max_steps, temperature, top_p, max_new_tokens]
        run_btn.click(start, inputs=inputs, outputs=outputs, api_name="start")
        prompt.submit(start, inputs=inputs, outputs=outputs, api_visibility="private")
        # Each agent step is its own request (fresh ZeroGPU token); the tick change chains the next one.
        adv = tick.change(advance, inputs=[state], outputs=outputs, api_name="advance", concurrency_limit=None)
        stop_btn.click(stop_run, inputs=[state], outputs=outputs, cancels=[adv], api_name="stop")
        gr.Examples(examples=[[e] for e in EXAMPLES], inputs=[prompt], outputs=outputs, fn=start,
                    cache_examples=False, run_on_click=True)

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)