"""Pre-bake station + cursed broadcasts into broadcasts_cache.json.

Generates POOL_SIZE wording variants for every spoken station, in every
language, plus the cursed signal. The Space loads this file at startup and
serves those broadcasts instantly (no model call), while each listener still
lands on one of several editions (variant % POOL_SIZE), so the per-listener
feel is kept. The encrypted number station and the hidden-band fragments are
fixed text already and are not baked; the operator console stays fully live.

Run once, locally (it is slow on CPU, a while). Re-run only when you change a
station persona, a prompt, POOL_SIZE, or the model:

    python prebake.py

The keys written here must match exactly what app.tune() looks up:
    station:{canonical}:{lang}:{pool}
    cursed:{CURSED_FREQUENCY}:{lang}:{pool}
"""

import json
import os
import time
from pathlib import Path

# Tell app.py not to spawn its background warm-up thread when we import it.
os.environ["PREBAKE"] = "1"

from app import POOL_SIZE, _mix_seed, _strip_leaks  # noqa: E402
from model import stream_broadcast  # noqa: E402
from stations import (  # noqa: E402
    CURSED_FREQUENCY,
    CURSED_SYSTEM_PROMPT,
    STATION_FREQUENCIES,
    get_station,
    localized_system_prompt,
    localized_user_prompt,
    station_seed,
    station_system_prompt,
)

OUT = Path(__file__).parent / "broadcasts_cache.json"
LANGS = ("es", "en", "fr")


def _generate(system_prompt: str, user_prompt: str, seed: int, max_tokens: int = 160) -> str:
    """One full broadcast, cleaned exactly like the live server caches it."""
    accumulated = ""
    for token in stream_broadcast(
        system_prompt, user_prompt=user_prompt, seed=seed, max_tokens=max_tokens
    ):
        accumulated += token
    return _strip_leaks(accumulated).strip()


def main():
    # Resume support: keep anything already baked so re-runs/interrupts are cheap.
    cache: dict[str, str] = {}
    if OUT.exists():
        try:
            cache = json.loads(OUT.read_text(encoding="utf-8"))
            print(f"resuming: {len(cache)} broadcasts already baked")
        except Exception:
            cache = {}

    jobs = []  # (key_prefix, system_prompt_fn_args)
    for freq in STATION_FREQUENCIES:
        station = get_station(freq)
        for lang in LANGS:
            jobs.append(
                (
                    f"station:{freq}:{lang}",
                    station_system_prompt(freq, lang),
                    localized_user_prompt(station, lang),
                    freq,
                )
            )
    for lang in LANGS:
        jobs.append(
            (
                f"cursed:{CURSED_FREQUENCY}:{lang}",
                localized_system_prompt(CURSED_SYSTEM_PROMPT, "es", lang),
                localized_user_prompt(None, lang),
                CURSED_FREQUENCY,
            )
        )

    total = len(jobs) * POOL_SIZE
    done = 0
    t0 = time.time()
    for prefix, sysp, usrp, freq in jobs:
        for pool in range(POOL_SIZE):
            done += 1
            key = f"{prefix}:{pool}"
            if cache.get(key):
                print(f"[{done}/{total}] {key}  (cached, skip)")
                continue
            text = _generate(sysp, usrp, _mix_seed(station_seed(freq), pool))
            if text:
                cache[key] = text
            # Write incrementally so an interrupt never loses progress.
            OUT.write_text(json.dumps(cache, ensure_ascii=False), encoding="utf-8")
            elapsed = time.time() - t0
            print(f"[{done}/{total}] {key}  ({len(text)} chars, {elapsed:.0f}s)")

    print(f"\nDone: {len(cache)} broadcasts -> {OUT} in {time.time() - t0:.0f}s")


if __name__ == "__main__":
    main()