#!/usr/bin/env python3
"""Generates transcript.jsonl -- the synthetic fixture the benchmark opens.

Deterministic (fixed seed), so runs draw byte-identical content.

Never a real transcript -- see AGENTS.md's ui-sandbox.sh, which this borrows its vocabulary
style from (headings, code fences, a table, a link) rather than reusing its Claude-Code JSONL
shape. This file's shape is the *app's own event model* instead: one JSON object per line,
matching what GET /sessions/{id}/transcript returns
(server/src/session/driver.rs is the source of truth for the field names).

    ./generate.py                      writes transcript.jsonl and bench1.png/bench2.png here

BACKLOG_COUNT events (seq 1..BACKLOG_COUNT) are the scrolled-back history the benchmark opens
with. A further STREAM_COUNT events (seq BACKLOG_COUNT+1..) are not part of the opening window;
both bench harnesses replay them at a fixed rate as the "streaming reply" phase, appended through
the same live path a real SSE reply arrives on. Keeping both halves in one file means one
generator and one seed to keep in sync, rather than two fixtures that can drift apart.
"""
import base64
import json
import random
import struct
import zlib
from pathlib import Path

SEED = 20260905
# The stress message below, two events, added to the 3,200 the backlog used to be. Both apps
# hardcode this to split the file by line index (`fixture.rs`'s `BACKLOG_COUNT`,
# `BenchFixture.kt`'s), so it moves in three places or none.
STRESS_EVENTS = 2
BACKLOG_COUNT = 3200 + STRESS_EVENTS
STREAM_COUNT = 400

# How many deltas fold into one markdown block of the streamed reply.
#
# **The blank line between them is the point** (2026-09-09). Until this run the streaming tail
# appended `paragraph(5) + " "` STREAM_COUNT times with no blank line anywhere, so all 400 deltas
# folded into a *single* 14,888-character block -- and a row re-shapes the block a delta lands in
# (`RowBlocks::apply_delta`), so every delta re-shaped the whole thing. That is quadratic in the
# reply's length, and it put 9.5ms of frame time on Iris's phone into a shape that does not occur:
# measured over 7,706 top-level blocks from 3,675 real assistant messages, block length is p50 147
# characters, p90 449, p99 836, largest 1,580, nothing above 4,000. `paragraph(5)` is ~35
# characters, so 4-12 of them per block lands in that range.
#
# The run-on version is kept, as STRESS_CHARS below -- Iris, 2026-09-09: "let's switch to new
# lines for the test, and also let's keep the single line around for stress".
DELTAS_PER_BLOCK = (4, 12)

# One deliberately pathological message in the backlog: a single markdown block with no blank line
# in it, the shape the streaming tail used to have. Sized just under `text_cap`'s MESSAGE_BYTES
# (16 KiB) so it is drawn in full rather than behind a "Show all" -- which makes it a genuine
# stress for one-shot shaping and puts a row right on the cap boundary, where nothing else is.
#
# It is *not* streamed: the repeated-reshape pathology needs a growing block, and that lives in
# the UI profiling rig's `what_reshaping_a_growing_message_costs`, where it can be
# iterated on in a second rather than in a two-minute phone run. Note that the cap would not save
# a real one anyway -- `row::build_row`'s `cap` is deliberately `false` for the live tail, because
# a row that grew while capped would appear to stop growing, so a streamed block's shaping cost
# has no ceiling.
STRESS_CHARS = 14_800
HERE = Path(__file__).resolve().parent / "assets"

random.seed(SEED)

LANGUAGES = ["rust", "kotlin", "python", "sh", "json", "toml"]

CODE_SNIPPETS = {
    "rust": '''fn fold_event(items: Vec<Item>, seq: u64) -> Vec<Item> {
    // a comment worth keeping: this is the fold the app's own screen runs
    let mut out = items;
    out.push(Item::new(seq));
    out
}''',
    "kotlin": '''fun foldEvent(items: List<TranscriptItem>, entry: SeqEvent): List<TranscriptItem> {
    // mirrors the server's own event model, one item per line
    return items + TranscriptItem.from(entry)
}''',
    "python": '''def render_report(frames, cpu_ms, rss_kb):
    # printed for a human to paste back, so every number carries its unit
    return f"{frames} frames, {cpu_ms}ms cpu, {rss_kb}kb peak rss"''',
    "sh": '''#!/bin/sh
# scripted scroll loop, the shape transcript-bench.sh drives on a phone
for i in $(seq 1 24); do
    ui-trace record --do "swipe 540 700 540 1600 200"
done''',
    "json": '{"seq": 1, "type": "status", "state": "running"}',
    "toml": '''[package]
name = "bench-fixture"
version = "0.1.0"''',
}

HEADINGS = [
    "## Plan",
    "## What changed",
    "## Why this approach",
    "### Open questions",
    "## Results",
]

WORDS = (
    "session render report frame budget scroll transcript fold event cache "
    "cursor probe stream backlog swipe fixture bench compose iris widget layout "
    "measure place draw tool call token context window anchor"
).split()


def paragraph(n=24):
    words = [random.choice(WORDS) for _ in range(n)]
    words[0] = words[0].capitalize()
    text = " ".join(words) + "."
    # Sprinkle markdown inline spans so the syntax highlighter/markdown parser sees a real mix.
    text = text.replace(" fold ", " **fold** ", 1)
    text = text.replace(" cursor ", " *cursor* ", 1)
    text = text.replace(" cache ", " `cache` ", 1)
    if "bench" in text:
        text = text.replace(
            " bench ", " [bench](https://example.com/bench) ", 1
        )
    return text


def make_png(rgb, size=8):
    """A tiny, valid PNG -- flat colour, no external dependency."""

    def chunk(tag, data):
        c = tag + data
        return struct.pack(">I", len(data)) + c + struct.pack(">I", zlib.crc32(c))

    sig = b"\x89PNG\r\n\x1a\n"
    ihdr = struct.pack(">IIBBBBB", size, size, 8, 2, 0, 0, 0)
    raw = b""
    for _ in range(size):
        raw += b"\x00" + bytes(rgb) * size
    idat = zlib.compress(raw)
    return sig + chunk(b"IHDR", ihdr) + chunk(b"IDAT", idat) + chunk(b"IEND", b"")


def main():
    HERE.mkdir(exist_ok=True)
    lines = []
    seq = 1
    ts = 1_788_000_000.0

    def emit(type_, **fields):
        nonlocal seq, ts
        obj = {"seq": seq, "ts": round(ts, 3), "type": type_}
        obj.update(fields)
        lines.append(json.dumps(obj, separators=(",", ":")))
        seq += 1
        ts += random.uniform(0.05, 2.0)

    emit("status", state="running")
    emit("settings", model="bench-model", permissionMode="auto")

    # The stress message, emitted first and with the random state put back afterwards, so that
    # adding it is **purely additive**: every turn the loop below generates is byte-identical to
    # what it generated before this existed, and only the streaming tail changed. That matters
    # because the fixture's opening view is what `phone_screen.rs`'s recorded gestures press --
    # `a_long_press_and_drag_selects_text` drives a real recording at (300, 1000) and fails the
    # moment different content lands under it. `ts` still shifts by the two draws, which nothing
    # reads.
    rng_state = random.getstate()
    emit(
        "userMessage",
        text="And the pathological one: a reply with no paragraph break in it.",
        id=None,
        attachments=[],
    )
    run_on = ""
    while len(run_on) < STRESS_CHARS:
        run_on += paragraph(5) + " "
    emit("assistantText", delta=run_on.rstrip())
    random.setstate(rng_state)

    image_refs = []
    turn = 0
    while seq <= BACKLOG_COUNT:
        turn += 1
        emit("userMessage", text=f"Turn {turn}: {paragraph(12)}", id=None, attachments=[])

        # A tool call with kilobyte-scale input/output every few turns.
        if turn % 3 == 0:
            tool_id = f"tool-{turn}"
            big_input = json.dumps({"path": f"/repo/file_{turn}.rs", "content": paragraph(400)})
            emit("toolStart", id=tool_id, tool="Edit", input=big_input)
            big_output = "\n".join(paragraph(60) for _ in range(20))
            emit("toolUpdate", id=tool_id, output=big_output[: len(big_output) // 2])
            emit("toolEnd", id=tool_id, output=big_output)

        # A reply: a heading, prose, a fenced block in a rotating language, a table, then deltas.
        emit("assistantText", delta=f"{random.choice(HEADINGS)}\n\n")
        emit("assistantText", delta=paragraph(30) + "\n\n")
        lang = LANGUAGES[turn % len(LANGUAGES)]
        emit("assistantText", delta=f"```{lang}\n{CODE_SNIPPETS[lang]}\n```\n\n")
        if turn % 5 == 0:
            emit(
                "assistantText",
                delta="| column | value |\n|---|---|\n| a | " + paragraph(3) + " |\n\n",
            )
        # A run of small deltas -- the shape a live reply actually streams in.
        for _ in range(random.randint(3, 8)):
            emit("assistantText", delta=paragraph(6) + " ")

        # A couple of images, base64 PNGs, the way a real transcript embeds a screenshot.
        if turn in (10, 40):
            ref = f"bench{len(image_refs) + 1}.png"
            image_refs.append(ref)
            emit("image", ref=ref, about=None)

        emit("usageDelta", tokens=random.randint(200, 4000), context=random.randint(2000, 180000))

        if turn % 15 == 0:
            emit(
                "compacted",
                preTokens=180000,
                postTokens=20000,
                trigger="auto",
            )

    # The streaming-phase tail: one reply built entirely from text deltas, the shape a bench
    # harness replays at a fixed events/sec through the live fold path -- with a blank line every
    # few deltas, so the block a delta lands in stays the size a real reply's blocks are. See
    # DELTAS_PER_BLOCK for what the alternative measured.
    emit("userMessage", text="One more, streamed live for the benchmark's timing phase.", id=None, attachments=[])
    until_break = random.randint(*DELTAS_PER_BLOCK)
    while seq <= BACKLOG_COUNT + STREAM_COUNT:
        until_break -= 1
        if until_break <= 0:
            # The break rides on the last delta of the block rather than being an event of its
            # own, so STREAM_COUNT still counts deltas a reader sees text arrive from.
            emit("assistantText", delta=paragraph(5) + "\n\n")
            until_break = random.randint(*DELTAS_PER_BLOCK)
        else:
            emit("assistantText", delta=paragraph(5) + " ")
    emit("status", state="idle")

    (HERE / "transcript.jsonl").write_text("\n".join(lines) + "\n")

    (HERE / "bench1.png").write_bytes(make_png((220, 90, 90)))
    (HERE / "bench2.png").write_bytes(make_png((90, 150, 220)))

    print(f"wrote {len(lines)} events ({BACKLOG_COUNT} backlog + {STREAM_COUNT} stream) to transcript.jsonl")


if __name__ == "__main__":
    main()