"""Record a short demo video of laya driving a real browser through jev-ultrafast.

    python apps/make_demo.py out_dir          (services: chromium :9222, laya systemone :8791, sglang Qwen :30000 for typed text)

Runs multi-step tasks on held-out sites (search, filters, sort, open a result) with frame recording, overlays goal / laya's decision / per-step latency, renders MP4 + GIF via ffmpeg.
"""
import base64, json, os, shutil, subprocess, sys, time
from pathlib import Path
sys.path.insert(0, "/home/ckl/projects/S/jev-ultrafast")
os.environ.update(BU_CDP_URL="http://127.0.0.1:9222", TYPESAFE_BASE_URL="http://127.0.0.1:8791", TYPESAFE_API_KEY="local", TEXT_MODEL_API_KEY="local",
                  TEXT_MODEL_BASE_URL="http://127.0.0.1:30000/v1", TEXT_MODEL="Qwen/Qwen3-8B-AWQ",
                  TEXT_MODEL_EXTRA_JSON='{"chat_template_kwargs": {"enable_thinking": false}}')   # typing tasks need the local text model
from PIL import Image, ImageDraw, ImageFont
from jev_ultrafast import Agent

TASKS = [   # multi-step tasks on held-out sites (suite C: none of these domains is in any training source)
    ("https://wordpress.org/plugins/", "Search the plugin directory for 'cache', show only commercial plugins, and open the 'Redis Object Cache' plugin."),
    ("https://packages.fedoraproject.org/", "Search for 'vim' and open the package 'vim-common'."),
    ("https://pkgs.racket-lang.org/", "Search packages for 'json' and open the package 'argo'."),
]
W, H = 1120, 780; BAR_T, BAR_B = 64, 96; FPS = 24
F_REG = "/usr/share/fonts/TTF/DejaVuSans.ttf"; F_BOLD = "/usr/share/fonts/TTF/DejaVuSans-Bold.ttf"; F_MONO = "/usr/share/fonts/TTF/DejaVuSansMono.ttf"
def font(n, bold=False): return ImageFont.truetype(F_BOLD if bold else F_REG, n)
def mono(n): return ImageFont.truetype(F_MONO, n)
BG, INK, MUTED, GREEN, BLUE = "#0f1115", "#e8e8e8", "#8b93a1", "#3ddc84", "#4f8cff"

def canvas(page_img=None):
    im = Image.new("RGB", (W, BAR_T + H + BAR_B), BG)
    if page_img is not None: im.paste(page_img.resize((W, H)), (0, BAR_T))
    return im

def frame(page_img, goal, step_txt, decision_txt, lat_txt, elapsed_txt, status=None):
    im = canvas(page_img); d = ImageDraw.Draw(im)
    d.text((16, 12), "laya-browser", font=font(22, True), fill=GREEN); d.text((190, 16), "System-1 decision head · 322M · local RTX 4070", font=font(15), fill=MUTED)
    d.text((16, 38), "goal: " + goal, font=font(16), fill=INK)
    y = BAR_T + H + 12
    d.text((16, y), step_txt, font=font(15, True), fill=MUTED)
    d.text((16, y + 24), decision_txt, font=mono(19), fill=INK)
    d.text((16, y + 54), lat_txt, font=mono(15), fill=BLUE); d.text((W - 200, y + 54), elapsed_txt, font=mono(15), fill=MUTED)
    if status: d.text((W - 200, y), status, font=font(17, True), fill=GREEN if status == "DONE" else "#ff6b6b")
    return im

def title_card(lines, sub):
    im = canvas(); d = ImageDraw.Draw(im); y = 260
    for i, l in enumerate(lines):
        d.text((60, y), l, font=font(40 if i == 0 else 26, i == 0), fill=GREEN if i == 0 else INK); y += 64 if i == 0 else 40
    y += 20
    for l in sub: d.text((60, y), l, font=font(20), fill=MUTED); y += 32
    return im

def run_task(url, goal, rec):
    frames, decisions = [], []
    with Agent(url, goal, record_dir=str(rec), screenshots=True) as agent:
        t0 = time.time(); page0 = agent.state["page"]
        frames.append((0, Image.open(rec / "000000.jpg").convert("RGB"), None))
        while agent.state["status"] not in ("done", "blocked") and len(agent.state["history"]) < 8:
            st = agent.command("tick")
            d = st["decisions"][-1] if st["decisions"] else None
            h = st["history"][-1] if st["history"] else None
            lat = d["latency_ms"] if d else 0
            label = (h["action"] if h and d and h.get("latency_ms") == lat else None)
            dec = {"op": d["operation"] if d else "?", "target": label or (d.get("target") if d else ""), "conf": d["confidence"] if d else 0, "lat": lat,
                   "elapsed": st["elapsed_ms"], "status": st["status"]}
            decisions.append(dec)
            shots = sorted(p for p in rec.glob("*.jpg") if p.stem != "000000")
            img = Image.open(shots[-1]).convert("RGB") if shots else frames[-1][1]
            frames.append((st["elapsed_ms"], img, dec))
        final = agent.state
    return frames, decisions, final

def main():
    out = Path(sys.argv[1]); shutil.rmtree(out, ignore_errors=True); (out / "frames").mkdir(parents=True)
    n = 0
    def emit(im, secs):
        nonlocal n
        for _ in range(int(secs * FPS)):
            im.save(out / "frames" / f"{n:06d}.png"); n += 1
    emit(title_card(["laya-browser", "one bidirectional encoder pass per step -> operation + target + calibrated confidence",
                     "no text generation, ~20-25 ms per decision on an RTX 4070 Ti SUPER, 1.5 GB VRAM", "multi-step tasks on sites that appear in no training source"],
                    ["fine-tuned from convaiinnovations/laya on Mind2Web + WebChain human trajectories + webgym + on-policy corrections",
                     "driving browser-use/jev-ultrafast through its TypeSafe-compatible /v1/systemone API"]), 3.5)
    totals = {"steps": 0, "lat": [], "wall": 0.0, "tasks_ok": 0}
    # warm-up pass (not recorded): the fast path compiles a kernel the first time it meets an input shape, which would
    # otherwise show up as 200-600 ms "decisions" in the video
    for i, (url, goal) in enumerate(TASKS):
        try:
            run_task(url, goal, out / f"warm{i}")
        except Exception as e:
            print("warm-up failed:", goal[:40], type(e).__name__)
        shutil.rmtree(out / f"warm{i}", ignore_errors=True)
    for i, (url, goal) in enumerate(TASKS):
        for attempt in range(3):          # live sites time out now and then: keep the first attempt that completes
            rec = out / f"rec{i}_{attempt}"; rec.mkdir()
            try:
                t = time.time(); frames, decisions, final = run_task(url, goal, rec); wall = time.time() - t
            except Exception as e:
                print("task failed:", goal[:50], type(e).__name__, str(e)[:60]); continue
            if final["status"] == "done":
                break
            print("not done, retrying:", goal[:50], final["status"])
        else:
            continue
        ok = final["status"] == "done"
        totals["steps"] += len(decisions); totals["lat"] += [d["lat"] for d in decisions]; totals["wall"] += wall; totals["tasks_ok"] += ok
        print(f"{goal[:50]:50s} steps={len(decisions)} status={final['status']} wall={wall:.1f}s decisions={[ (d['op'], d['lat']) for d in decisions]}", flush=True)
        # observed page, then each decision: show the page it decided on with the decision, then the resulting page
        emit(frame(frames[0][1], goal, "step 0 · observe page", "laya reads the element table ...", "", "t = 0 ms"), 1.2)
        for k, (ms, img, dec) in enumerate(frames[1:], 1):
            prev = frames[k - 1][1]
            dtxt = f"{dec['op']}  ->  {str(dec['target'])[:48]}" if dec["op"] not in ("DONE", "BLOCKED") else dec["op"]
            emit(frame(prev, goal, f"step {k} · decide", dtxt, f"laya: {dec['lat']} ms  (confidence {dec['conf']:.2f})", f"t = {dec['elapsed']} ms"), 1.3)
            if dec["op"] not in ("DONE", "BLOCKED"):
                emit(frame(img, goal, f"step {k} · executed", dtxt, f"laya: {dec['lat']} ms", f"t = {dec['elapsed']} ms"), 1.0)
            else:
                emit(frame(img, goal, f"step {k}", dtxt, f"laya: {dec['lat']} ms", f"t = {dec['elapsed']} ms", status=dec["op"]), 1.6)
    lat = sorted(totals["lat"]); med = lat[len(lat) // 2] if lat else 0
    emit(title_card(["that's laya as a browser agent's System 1", f"{totals['tasks_ok']}/{len(TASKS)} tasks · {totals['steps']} decisions · median {med} ms per decision",
                     f"total wall time {totals['wall']:.1f} s including page loads"],
                    ["model + code + numbers: huggingface.co/cklxx/laya-browser", "TileLang fused kernels + CUDA graphs: github.com/NandhaKishorM/laya/pull/25"]), 4.0)
    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-framerate", str(FPS), "-i", str(out / "frames" / "%06d.png"), "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "23", str(out / "laya_browser_demo.mp4")], check=True)
    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(out / "laya_browser_demo.mp4"), "-vf", "fps=8,scale=700:-1:flags=lanczos,split[s0][s1];[s0]palettegen=max_colors=128[p];[s1][p]paletteuse=dither=bayer", str(out / "laya_browser_demo.gif")], check=True)
    print("video:", out / "laya_browser_demo.mp4", f"{(out / 'laya_browser_demo.mp4').stat().st_size/1e6:.1f} MB", "| gif:", f"{(out / 'laya_browser_demo.gif').stat().st_size/1e6:.1f} MB", "| frames", n, f"({n/FPS:.0f}s)")

if __name__ == "__main__":
    main()