"""Standalone web UI for LTX-2.3 video generation.

Designed to be distributed inside a HuggingFace model repo.  Place this file
anywhere in the downloaded repo directory (or a sibling directory) and run:

    python webapp_standalone.py
    python webapp_standalone.py --port 8080
    python webapp_standalone.py --compare-dir /path/to/second/model
    python webapp_standalone.py --model-name "LTX-2.3 24 GB RAM"

The script auto-detects its own directory as the primary model.  A second
model can be supplied via --compare-dir for A/B comparison.

Requirements:
    pip install flask
    pip install mlx mlx-lm ltx-core-mlx ltx-pipelines-mlx
"""

import argparse
import json
import subprocess
import sys
import threading
import time
import uuid
from collections import defaultdict
from pathlib import Path

from flask import Flask, Response, jsonify, request, send_file

# ---------------------------------------------------------------------------
# CLI args — parsed at import time so the constants below can reference them
# ---------------------------------------------------------------------------

def _build_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="LTX-2.3 standalone video-generation web UI"
    )
    p.add_argument("--port", type=int, default=7860,
                   help="Port to listen on (default: 7860)")
    p.add_argument("--compare-dir", type=str, default=None,
                   help="Optional path to a second model directory for A/B comparison")
    p.add_argument("--model-name", type=str, default=None,
                   help="Display name for the primary model (default: directory name)")
    return p


# Parse only our own args; anything unrecognised is left alone so Flask's own
# dev-server reloader doesn't choke on our flags.
_parser = _build_arg_parser()
_args, _unknown = _parser.parse_known_args()

# ---------------------------------------------------------------------------
# Model discovery
# ---------------------------------------------------------------------------

# The primary model IS the directory that contains this script — i.e. the
# downloaded HuggingFace repo root.
PRIMARY_DIR = Path(__file__).parent.resolve()
PRIMARY_NAME = _args.model_name or PRIMARY_DIR.name

COMPARE_DIR: Path | None = Path(_args.compare_dir).resolve() if _args.compare_dir else None
COMPARE_NAME: str | None = COMPARE_DIR.name if COMPARE_DIR else None

# Required files that signal a valid, ready model directory.
_REQUIRED_FILES = [
    "transformer-distilled.safetensors",
    "connector.safetensors",
    "vae_decoder.safetensors",
    "audio_vae.safetensors",
    "vocoder.safetensors",
]


def _model_ready(path: Path) -> bool:
    return path.is_dir() and all((path / f).exists() for f in _REQUIRED_FILES)


def _model_missing_files(path: Path) -> list[str]:
    return [f for f in _REQUIRED_FILES if not (path / f).exists()]


# Build the static model list once at startup.
MODELS: list[dict] = []

_primary_ok = _model_ready(PRIMARY_DIR)
MODELS.append({
    "id":    "primary",
    "label": PRIMARY_NAME,
    "dir":   str(PRIMARY_DIR),
    "ready": _primary_ok,
    "missing": _model_missing_files(PRIMARY_DIR) if not _primary_ok else [],
})

if COMPARE_DIR is not None:
    _compare_ok = _model_ready(COMPARE_DIR)
    MODELS.append({
        "id":    "compare",
        "label": COMPARE_NAME,
        "dir":   str(COMPARE_DIR),
        "ready": _compare_ok,
        "missing": _model_missing_files(COMPARE_DIR) if not _compare_ok else [],
    })

# Convenience lookup: id → dir
MODEL_DIRS: dict[str, str] = {m["id"]: m["dir"] for m in MODELS}

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------

# Videos are saved alongside this script (which lives in the model repo).
RESULTS_DIR = PRIMARY_DIR / "webapp_videos"
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

# generate_ltx.py is located relative to the RAM/RUN working tree.  We find
# it by searching upward from this file, then falling back to a path the user
# can override via the GENERATE_SCRIPT env var.
import os as _os

def _find_generate_script() -> Path:
    env_override = _os.environ.get("GENERATE_SCRIPT")
    if env_override:
        return Path(env_override)
    # Walk up looking for experiments/flux_phase1/generate_ltx.py
    cur = Path(__file__).parent
    for _ in range(6):
        candidate = cur / "experiments" / "flux_phase1" / "generate_ltx.py"
        if candidate.exists():
            return candidate
        cur = cur.parent
    # Last resort: assume this script is inside RAM/RUN/results/<something>/
    # so climb two levels to RAM/RUN/
    return Path(__file__).parent.parent.parent / "experiments" / "flux_phase1" / "generate_ltx.py"


GENERATE_SCRIPT = _find_generate_script()

# ---------------------------------------------------------------------------
# Job state
# ---------------------------------------------------------------------------

# job_id → {status, log_lines, video_path, started, finished, pid, params}
JOBS: dict = {}
JOBS_LOCK = threading.Lock()

# ---------------------------------------------------------------------------
# apply_mixed_precision_quantization
# (kept here so the script is self-contained; also used by generate_ltx.py
# which is invoked as a subprocess — but having it here lets us surface the
# logic for anyone reading this file)
# ---------------------------------------------------------------------------

def apply_mixed_precision_quantization(model, weights, group_size: int = 64) -> None:
    """Per-layer mixed-precision quantization from a weight dict.

    Unlike ltx_core_mlx's apply_quantization (which uses a single detected
    bit width for all layers), this version detects each layer's bits from
    its packed weight shape and applies nn.quantize once per unique bit width.
    """
    import mlx.nn as nn

    layer_bits: dict[str, int] = {}
    for key in weights:
        if not key.endswith(".scales"):
            continue
        layer = key[: -len(".scales")]
        w_key = layer + ".weight"
        if w_key not in weights:
            continue
        w_cols = weights[w_key].shape[-1]
        s_cols = weights[key].shape[-1]
        bits = round(w_cols * 32 / (s_cols * group_size))
        if bits in (2, 3, 4, 5, 6, 8):
            layer_bits[layer] = bits

    if not layer_bits:
        return

    bits_to_layers: dict[int, set] = defaultdict(set)
    for layer, b in layer_bits.items():
        bits_to_layers[b].add(layer)

    for bits, layers in sorted(bits_to_layers.items()):
        def _predicate(path: str, module, _layers=layers) -> bool:
            return path in _layers and isinstance(module, nn.Linear)
        nn.quantize(model, group_size=group_size, bits=bits, class_predicate=_predicate)

    total = sum(len(v) for v in bits_to_layers.values())
    dist = {b: len(v) for b, v in sorted(bits_to_layers.items())}
    print(f"  Mixed-precision quantization: {total} layers — {dist}", flush=True)

# ---------------------------------------------------------------------------
# Job runner
# ---------------------------------------------------------------------------

def _run_job(job_id: str, cmd: list[str], video_path: Path, cwd: str):
    with JOBS_LOCK:
        JOBS[job_id]["status"] = "running"

    try:
        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            cwd=cwd,
        )
        with JOBS_LOCK:
            JOBS[job_id]["pid"] = proc.pid

        for line in proc.stdout:
            line = line.rstrip("\n")
            with JOBS_LOCK:
                JOBS[job_id]["log_lines"].append(line)

        proc.wait()
        success = proc.returncode == 0 and video_path.exists()
        with JOBS_LOCK:
            JOBS[job_id]["status"] = "done" if success else "error"
            JOBS[job_id]["finished"] = time.time()
            if success:
                JOBS[job_id]["video_path"] = str(video_path)
    except Exception as exc:
        with JOBS_LOCK:
            JOBS[job_id]["log_lines"].append(f"[webapp error] {exc}")
            JOBS[job_id]["status"] = "error"
            JOBS[job_id]["finished"] = time.time()

# ---------------------------------------------------------------------------
# Flask app
# ---------------------------------------------------------------------------

app = Flask(__name__)


@app.post("/generate")
def generate():
    data = request.get_json(force=True)
    prompt = data.get("prompt", "").strip()
    if not prompt:
        return jsonify(error="prompt required"), 400

    model_id = data.get("model", MODELS[0]["id"])
    if model_id not in MODEL_DIRS:
        return jsonify(error=f"unknown model: {model_id}"), 400
    model_dir = MODEL_DIRS[model_id]
    if not Path(model_dir).exists():
        return jsonify(error=f"model directory not found: {model_dir}"), 400

    if not _model_ready(Path(model_dir)):
        missing = _model_missing_files(Path(model_dir))
        return jsonify(error=f"model not ready, missing: {missing}"), 400

    height     = int(data.get("height", 480))
    width      = int(data.get("width", 704))
    num_frames = int(data.get("num_frames", 65))
    frame_rate = float(data.get("frame_rate", 24.0))
    seed       = int(data.get("seed", 42))
    stage1     = data.get("stage1_steps")
    stage2     = data.get("stage2_steps")

    job_id = uuid.uuid4().hex[:8]
    video_path = RESULTS_DIR / f"gen_{job_id}.mp4"

    if not GENERATE_SCRIPT.exists():
        return jsonify(error=(
            f"generate_ltx.py not found at {GENERATE_SCRIPT}. "
            "Set the GENERATE_SCRIPT env var to its absolute path."
        )), 500

    # Determine the cwd for the subprocess.  generate_ltx.py expects to be
    # run from the RAM/RUN/ directory so its relative imports resolve.
    script_cwd = str(GENERATE_SCRIPT.parent.parent.parent)

    cmd = [
        sys.executable,
        str(GENERATE_SCRIPT),
        "--model-dir", model_dir,
        "--prompt", prompt,
        "--output", str(video_path),
        "--height", str(height),
        "--width",  str(width),
        "--num-frames", str(num_frames),
        "--frame-rate", str(frame_rate),
        "--seed", str(seed),
    ]
    if stage1:
        cmd += ["--stage1-steps", str(stage1)]
    if stage2:
        cmd += ["--stage2-steps", str(stage2)]

    # Derive a friendly label for the model in job params
    model_label = next((m["label"] for m in MODELS if m["id"] == model_id), model_id)

    with JOBS_LOCK:
        JOBS[job_id] = {
            "status": "queued",
            "log_lines": [],
            "video_path": None,
            "started": time.time(),
            "finished": None,
            "pid": None,
            "params": {
                "prompt": prompt,
                "model": model_id,
                "model_label": model_label,
                "height": height,
                "width": width,
                "num_frames": num_frames,
                "frame_rate": frame_rate,
                "seed": seed,
            },
        }

    t = threading.Thread(
        target=_run_job, args=(job_id, cmd, video_path, script_cwd), daemon=True
    )
    t.start()
    return jsonify(job_id=job_id)


@app.get("/stream/<job_id>")
def stream(job_id: str):
    """SSE live log stream for a running job."""
    if job_id not in JOBS:
        return jsonify(error="not found"), 404

    def generate_events():
        sent = 0
        while True:
            with JOBS_LOCK:
                lines  = JOBS[job_id]["log_lines"]
                status = JOBS[job_id]["status"]
                new_lines = lines[sent:]
                sent += len(new_lines)

            for line in new_lines:
                yield f"data: {json.dumps({'line': line})}\n\n"

            if status in ("done", "error") and not new_lines:
                with JOBS_LOCK:
                    final_status = JOBS[job_id]["status"]
                    video        = JOBS[job_id]["video_path"]
                yield f"data: {json.dumps({'done': True, 'status': final_status, 'video': video})}\n\n"
                return

            time.sleep(0.25)

    return Response(
        generate_events(),
        mimetype="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


@app.get("/video/<job_id>")
def video(job_id: str):
    with JOBS_LOCK:
        job = JOBS.get(job_id)
    if not job or not job["video_path"]:
        return jsonify(error="not found"), 404
    p = Path(job["video_path"])
    if not p.exists():
        return jsonify(error="file missing"), 404
    return send_file(str(p), mimetype="video/mp4", conditional=True)


@app.get("/models")
def list_models():
    """Return the static model list (no polling needed — models are local)."""
    return jsonify([
        {
            "id":      m["id"],
            "label":   m["label"],
            "dir":     m["dir"],
            "ready":   m["ready"],
            "missing": m["missing"],
        }
        for m in MODELS
    ])


@app.get("/jobs")
def list_jobs():
    with JOBS_LOCK:
        out = []
        for jid, j in reversed(list(JOBS.items())):
            out.append({
                "id":        jid,
                "status":    j["status"],
                "params":    j["params"],
                "started":   j["started"],
                "finished":  j["finished"],
                "has_video": bool(j["video_path"]),
            })
    return jsonify(out)


# ---------------------------------------------------------------------------
# HTML — single-file UI
# ---------------------------------------------------------------------------

def _build_html(models: list[dict]) -> str:
    # Build the model selector: single static label if one model, <select> if two.
    single_model = len(models) == 1

    if single_model:
        m = models[0]
        model_block = f"""
  <div>
    <div class="section-title">Model</div>
    <div id="model-display" style="
        padding:10px 12px;
        background:var(--bg);
        border:1px solid var(--border);
        border-radius:8px;
        font-size:.9rem;
        color:var(--text);
    ">{m['label']}</div>
    <input type="hidden" id="model" value="{m['id']}">
    <div id="model-note" style="font-size:.72rem;color:var(--muted);margin-top:5px;min-height:1.2em"></div>
  </div>"""
    else:
        options = "\n    ".join(
            f'<option value="{m["id"]}" {"disabled" if not m["ready"] else ""}>'
            f'{m["label"]}{"  (not ready)" if not m["ready"] else ""}'
            f'</option>'
            for m in models
        )
        # Default selection: first ready model
        default_id = next((m["id"] for m in models if m["ready"]), models[0]["id"])
        model_block = f"""
  <div>
    <div class="section-title">Model</div>
    <select id="model" onchange="updateModelNote()">
    {options}
    </select>
    <div id="model-note" style="font-size:.72rem;color:var(--muted);margin-top:5px;min-height:1.2em"></div>
  </div>"""

    # Startup warning if primary model is not ready
    startup_warn = ""
    if not models[0]["ready"]:
        missing_list = ", ".join(models[0]["missing"])
        startup_warn = f"""
<div style="
    grid-column:1/-1;
    background:#2a1a0a;
    border-bottom:1px solid #5a3a0a;
    padding:12px 28px;
    font-size:.82rem;
    color:#fbbf24;
">
  Model directory is missing required files: <code>{missing_list}</code>.
  Run <code>reformat_ltx_for_pipeline.py</code> first.
</div>"""

    # Models JSON for JS
    models_json = json.dumps([{"id": m["id"], "ready": m["ready"], "label": m["label"]} for m in models])

    return f"""<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>{models[0]['label']} — Video Generator</title>
<style>
  *, *::before, *::after {{ box-sizing: border-box; margin: 0; padding: 0; }}

  :root {{
    --bg:      #0f0f13;
    --surface: #1a1a22;
    --border:  #2e2e3a;
    --accent:  #7c6af7;
    --accent2: #a78bfa;
    --text:    #e2e2f0;
    --muted:   #6b6b82;
    --green:   #34d399;
    --red:     #f87171;
    --yellow:  #fbbf24;
  }}

  body {{
    background: var(--bg);
    color: var(--text);
    font-family: system-ui, -apple-system, sans-serif;
    min-height: 100vh;
    display: grid;
    grid-template-columns: 380px 1fr;
    grid-template-rows: auto auto 1fr;
    gap: 0;
  }}

  header {{
    grid-column: 1 / -1;
    padding: 18px 28px;
    border-bottom: 1px solid var(--border);
    display: flex;
    align-items: center;
    gap: 12px;
  }}
  header h1 {{ font-size: 1.1rem; font-weight: 600; letter-spacing: .02em; }}
  header .badge {{
    font-size: .7rem; background: var(--accent); color: #fff;
    padding: 2px 8px; border-radius: 99px; font-weight: 600;
  }}

  .warn-banner {{ grid-column: 1 / -1; }}

  .sidebar {{
    grid-column: 1;
    border-right: 1px solid var(--border);
    padding: 24px 20px;
    display: flex;
    flex-direction: column;
    gap: 18px;
    overflow-y: auto;
  }}

  .main {{
    grid-column: 2;
    padding: 24px 28px;
    display: flex;
    flex-direction: column;
    gap: 20px;
    overflow-y: auto;
  }}

  label {{ font-size: .8rem; color: var(--muted); display: block; margin-bottom: 5px; }}

  textarea, input, select {{
    width: 100%;
    background: var(--bg);
    border: 1px solid var(--border);
    border-radius: 8px;
    color: var(--text);
    padding: 10px 12px;
    font-size: .9rem;
    font-family: inherit;
    outline: none;
    transition: border-color .15s;
  }}
  textarea:focus, input:focus, select:focus {{ border-color: var(--accent); }}
  textarea {{ resize: vertical; min-height: 90px; }}

  .row {{ display: grid; grid-template-columns: 1fr 1fr; gap: 12px; }}

  button {{
    width: 100%;
    padding: 12px;
    border: none;
    border-radius: 8px;
    background: var(--accent);
    color: #fff;
    font-size: .95rem;
    font-weight: 600;
    cursor: pointer;
    transition: opacity .15s, background .15s;
  }}
  button:hover {{ opacity: .9; }}
  button:disabled {{ background: var(--border); color: var(--muted); cursor: not-allowed; opacity: 1; }}

  .section-title {{
    font-size: .7rem;
    font-weight: 700;
    letter-spacing: .1em;
    text-transform: uppercase;
    color: var(--muted);
    margin-bottom: -8px;
  }}

  .presets {{ display: flex; gap: 6px; flex-wrap: wrap; }}
  .preset {{
    font-size: .78rem; padding: 4px 10px; border-radius: 6px;
    border: 1px solid var(--border); background: var(--surface);
    cursor: pointer; color: var(--text); transition: border-color .15s;
    white-space: nowrap;
  }}
  .preset:hover, .preset.active {{ border-color: var(--accent); color: var(--accent2); }}

  .status-pill {{
    display: inline-flex; align-items: center; gap: 6px;
    font-size: .8rem; padding: 3px 10px; border-radius: 99px;
    font-weight: 600;
  }}
  .status-pill.queued  {{ background: #2a2a1a; color: var(--yellow); }}
  .status-pill.running {{ background: #1a1a2e; color: var(--accent2); }}
  .status-pill.done    {{ background: #0d2a1e; color: var(--green); }}
  .status-pill.error   {{ background: #2a0f0f; color: var(--red); }}
  .dot {{ width: 7px; height: 7px; border-radius: 50%; background: currentColor; }}
  .dot.pulse {{ animation: pulse 1s infinite; }}
  @keyframes pulse {{ 0%,100%{{opacity:1}} 50%{{opacity:.3}} }}

  .log-wrap {{
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 10px;
    overflow: hidden;
    flex: 1;
    min-height: 200px;
    display: flex;
    flex-direction: column;
  }}
  .log-header {{
    padding: 10px 14px;
    border-bottom: 1px solid var(--border);
    font-size: .78rem;
    color: var(--muted);
    display: flex;
    align-items: center;
    gap: 8px;
  }}
  .log-body {{
    flex: 1;
    overflow-y: auto;
    padding: 12px 14px;
    font-family: 'SF Mono', 'Fira Mono', monospace;
    font-size: .78rem;
    line-height: 1.6;
    color: #b0b0c8;
    white-space: pre-wrap;
    word-break: break-all;
    max-height: 300px;
  }}
  .log-body:empty::before {{ content: 'Waiting for output\2026'; color: var(--muted); }}

  .video-wrap {{
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 10px;
    overflow: hidden;
  }}
  .video-wrap video {{
    width: 100%;
    display: block;
    background: #000;
    max-height: 480px;
  }}
  .video-placeholder {{
    height: 200px;
    display: flex;
    align-items: center;
    justify-content: center;
    color: var(--muted);
    font-size: .85rem;
  }}

  .history-item {{
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 8px;
    padding: 10px 12px;
    font-size: .82rem;
    display: flex;
    flex-direction: column;
    gap: 4px;
    cursor: pointer;
    transition: border-color .15s;
    margin-bottom: 8px;
  }}
  .history-item:hover {{ border-color: var(--accent); }}
  .history-item .prompt {{ color: var(--text); font-weight: 500; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }}
  .history-item .meta {{ color: var(--muted); font-size: .75rem; }}

  .empty-state {{ color: var(--muted); font-size: .85rem; text-align: center; padding: 20px 0; }}

  code {{
    font-family: 'SF Mono', 'Fira Mono', monospace;
    font-size: .85em;
    background: #1e1e2a;
    padding: 1px 5px;
    border-radius: 4px;
  }}
</style>
</head>
<body>

<header>
  <h1>{models[0]['label']}</h1>
  <span class="badge">RAM Mixed-Precision</span>
</header>

{startup_warn}

<aside class="sidebar">
  {model_block}

  <div>
    <label for="prompt">Prompt</label>
    <textarea id="prompt" rows="4" placeholder="Describe the video you want to generate…">A serene mountain lake at sunrise, mist over calm water, pine trees reflected</textarea>
  </div>

  <div>
    <div class="section-title" style="margin-bottom:10px">Resolution presets</div>
    <div class="presets">
      <div class="preset" data-h="256"  data-w="256"  data-f="33"  title="256×256, 33 frames">tiny</div>
      <div class="preset active" data-h="480"  data-w="704"  data-f="65"  title="480×704, 65 frames">480p·65f</div>
      <div class="preset" data-h="480"  data-w="704"  data-f="97"  title="480×704, 97 frames">480p·97f</div>
      <div class="preset" data-h="720"  data-w="1280" data-f="65"  title="720×1280, 65 frames">720p·65f</div>
    </div>
  </div>

  <div class="row">
    <div>
      <label for="height">Height</label>
      <input type="number" id="height" value="480" step="32" min="128">
    </div>
    <div>
      <label for="width">Width</label>
      <input type="number" id="width" value="704" step="32" min="128">
    </div>
  </div>

  <input type="hidden" id="num_frames" value="65">

  <div class="row">
    <div>
      <div style="display:flex;justify-content:space-between;align-items:baseline;margin-bottom:5px">
        <label for="duration_s" style="margin:0">Duration</label>
        <span id="frames-display" style="font-size:.75rem;color:var(--muted)">=&thinsp;65 frames</span>
      </div>
      <div style="display:flex;gap:6px;align-items:center">
        <input type="number" id="duration_s" value="2.7" step="0.5" min="0.5" oninput="updateFrames()">
        <span style="color:var(--muted);font-size:.85rem;white-space:nowrap;padding-right:4px">s</span>
      </div>
    </div>
    <div>
      <label for="frame_rate">FPS</label>
      <input type="number" id="frame_rate" value="24" step="1" min="8" max="60" oninput="updateFrames()">
    </div>
  </div>

  <div class="row">
    <div>
      <label for="seed">Seed</label>
      <input type="number" id="seed" value="42">
    </div>
    <div style="display:flex;flex-direction:column;justify-content:flex-end">
      <button type="button" onclick="randomSeed()" style="padding:10px;font-size:.8rem;background:var(--surface);color:var(--text);border:1px solid var(--border)">\U0001f3b2 Random</button>
    </div>
  </div>

  <button id="btn-generate" onclick="startGeneration()">Generate</button>
</aside>

<main class="main">
  <div id="status-row" style="display:none;align-items:center;gap:12px">
    <span id="status-pill" class="status-pill queued"><span class="dot"></span> queued</span>
    <span id="status-info" style="font-size:.82rem;color:var(--muted)"></span>
  </div>

  <div id="video-section">
    <div class="video-wrap">
      <div class="video-placeholder" id="video-placeholder">Video will appear here after generation</div>
      <video id="video-el" controls style="display:none" playsinline></video>
    </div>
  </div>

  <div class="log-wrap">
    <div class="log-header">
      <span>Output log</span>
      <span id="log-status" style="margin-left:auto"></span>
    </div>
    <div class="log-body" id="log-body"></div>
  </div>

  <div>
    <div class="section-title" style="margin-bottom:12px">Recent generations</div>
    <div id="history"></div>
  </div>
</main>

<script>
// Static model list injected server-side — no polling needed.
const MODELS = {models_json};
const singleModel = MODELS.length === 1;

let currentJobId = null;
let currentESS   = null;

// ── Model note ────────────────────────────────────────────────────────────────

function updateModelNote() {{
  if (singleModel) return;
  const sel   = document.getElementById('model');
  const m     = MODELS.find(x => x.id === sel.value);
  const noteEl = document.getElementById('model-note');
  noteEl.textContent = (m && !m.ready) ? '⚠️ Model not ready — missing required files.' : '';
}}

// Run once on load if using <select>
if (!singleModel) updateModelNote();

// ── Helpers ───────────────────────────────────────────────────────────────────

function randomSeed() {{
  document.getElementById('seed').value = Math.floor(Math.random() * 2**31);
}}

function updateFrames() {{
  const secs = parseFloat(document.getElementById('duration_s').value) || 2.7;
  const fps  = parseFloat(document.getElementById('frame_rate').value) || 24;
  // LTX requires 32k+1 frames (k ≥ 1), minimum 33
  const k      = Math.max(1, Math.round((secs * fps - 1) / 32));
  const frames = k * 32 + 1;
  document.getElementById('num_frames').value = frames;
  document.getElementById('frames-display').textContent = `= ${{frames}} frames`;
}}

function syncDurationFromFrames(frames) {{
  const fps = parseFloat(document.getElementById('frame_rate').value) || 24;
  document.getElementById('duration_s').value = (frames / fps).toFixed(1);
  document.getElementById('num_frames').value  = frames;
  document.getElementById('frames-display').textContent = `= ${{frames}} frames`;
}}

document.querySelectorAll('.preset').forEach(el => {{
  el.addEventListener('click', () => {{
    document.querySelectorAll('.preset').forEach(p => p.classList.remove('active'));
    el.classList.add('active');
    document.getElementById('height').value = el.dataset.h;
    document.getElementById('width').value  = el.dataset.w;
    syncDurationFromFrames(parseInt(el.dataset.f));
  }});
}});

// ── Generation ────────────────────────────────────────────────────────────────

async function startGeneration() {{
  const prompt = document.getElementById('prompt').value.trim();
  if (!prompt) {{ alert('Enter a prompt first.'); return; }}

  const btn = document.getElementById('btn-generate');
  btn.disabled = true;

  document.getElementById('log-body').textContent = '';
  document.getElementById('video-el').style.display = 'none';
  document.getElementById('video-placeholder').style.display = 'flex';
  document.getElementById('status-row').style.display = 'flex';
  setStatus('queued');

  if (currentESS) {{ currentESS.close(); currentESS = null; }}

  const body = {{
    prompt,
    model:      document.getElementById('model').value,
    height:     parseInt(document.getElementById('height').value),
    width:      parseInt(document.getElementById('width').value),
    num_frames: parseInt(document.getElementById('num_frames').value),
    frame_rate: parseFloat(document.getElementById('frame_rate').value),
    seed:       parseInt(document.getElementById('seed').value),
  }};

  const res = await fetch('/generate', {{
    method: 'POST',
    headers: {{'Content-Type': 'application/json'}},
    body: JSON.stringify(body),
  }});
  if (!res.ok) {{
    const err = await res.json();
    alert('Error: ' + (err.error || res.statusText));
    btn.disabled = false;
    return;
  }}
  const {{ job_id }} = await res.json();
  currentJobId = job_id;

  document.getElementById('status-info').textContent = `job ${{job_id}}`;
  setStatus('running');

  const sse = new EventSource(`/stream/${{job_id}}`);
  currentESS = sse;
  const logEl = document.getElementById('log-body');

  sse.onmessage = (e) => {{
    const data = JSON.parse(e.data);
    if (data.line !== undefined) {{
      logEl.textContent += data.line + '\\n';
      logEl.scrollTop = logEl.scrollHeight;
    }}
    if (data.done) {{
      sse.close();
      currentESS = null;
      btn.disabled = false;
      setStatus(data.status);
      if (data.status === 'done' && data.video) showVideo(job_id);
      refreshHistory();
    }}
  }};
  sse.onerror = () => {{
    sse.close();
    currentESS = null;
    btn.disabled = false;
  }};
}}

function setStatus(s) {{
  const pill = document.getElementById('status-pill');
  pill.className = `status-pill ${{s}}`;
  const dot = pill.querySelector('.dot');
  dot.className = 'dot' + (s === 'running' ? ' pulse' : '');
  dot.nextSibling.textContent = ' ' + s;
}}

function showVideo(job_id) {{
  const el = document.getElementById('video-el');
  el.src = `/video/${{job_id}}`;
  el.style.display = 'block';
  document.getElementById('video-placeholder').style.display = 'none';
  el.load();
  el.play().catch(() => {{}});
}}

// ── History ───────────────────────────────────────────────────────────────────

async function refreshHistory() {{
  const res = await fetch('/jobs');
  if (!res.ok) return;
  const jobs = await res.json();
  const el   = document.getElementById('history');
  if (!jobs.length) {{
    el.innerHTML = '<div class="empty-state">No generations yet</div>';
    return;
  }}
  el.innerHTML = jobs.slice(0, 10).map(j => {{
    const ago      = Math.round((Date.now() / 1000 - j.started) / 60);
    const duration = j.finished ? `${{Math.round(j.finished - j.started)}}s` : '…';
    const modelLbl = j.params.model_label || j.params.model;
    return `<div class="history-item" onclick="loadJob('${{j.id}}')">
      <div class="prompt">${{escHtml(j.params.prompt)}}</div>
      <div class="meta">${{escHtml(modelLbl)}} · ${{j.params.height}}×${{j.params.width}} · ${{j.params.num_frames}}f · seed ${{j.params.seed}} · ${{duration}} · ${{ago}}m ago</div>
    </div>`;
  }}).join('');
}}

async function loadJob(job_id) {{
  const res  = await fetch('/jobs');
  const jobs = await res.json();
  const j    = jobs.find(x => x.id === job_id);
  if (!j) return;

  if (!singleModel) document.getElementById('model').value = j.params.model;
  document.getElementById('prompt').value     = j.params.prompt;
  document.getElementById('height').value     = j.params.height;
  document.getElementById('width').value      = j.params.width;
  document.getElementById('frame_rate').value = j.params.frame_rate;
  document.getElementById('seed').value       = j.params.seed;
  syncDurationFromFrames(j.params.num_frames);

  if (j.has_video) {{
    document.getElementById('video-el').style.display  = 'none';
    document.getElementById('video-placeholder').style.display = 'flex';
    showVideo(job_id);
    document.getElementById('status-row').style.display = 'flex';
    setStatus(j.status);
    document.getElementById('status-info').textContent = `job ${{job_id}}`;
  }}
}}

function escHtml(s) {{
  return s.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;').replace(/"/g,'&quot;');
}}

// ── Init ──────────────────────────────────────────────────────────────────────
refreshHistory();
</script>
</body>
</html>"""


@app.get("/")
def index():
    return _build_html(MODELS)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import webbrowser
    import threading as _threading

    port = _args.port

    # Print startup summary
    print(f"\nLTX-2.3 Standalone Web UI")
    print(f"  Primary model : {PRIMARY_NAME}")
    print(f"  Directory     : {PRIMARY_DIR}")
    print(f"  Ready         : {_model_ready(PRIMARY_DIR)}")
    if COMPARE_DIR:
        print(f"  Compare model : {COMPARE_NAME}")
        print(f"  Compare dir   : {COMPARE_DIR}")
        print(f"  Compare ready : {_model_ready(COMPARE_DIR)}")
    print(f"  Generate script: {GENERATE_SCRIPT} ({'found' if GENERATE_SCRIPT.exists() else 'NOT FOUND — set GENERATE_SCRIPT env var'})")
    print(f"  Videos saved to: {RESULTS_DIR}")
    print(f"\n  http://localhost:{port}\n")

    def _open():
        time.sleep(1.0)
        webbrowser.open(f"http://localhost:{port}")

    _threading.Thread(target=_open, daemon=True).start()
    app.run(host="0.0.0.0", port=port, debug=False, threaded=True)