"""Generate and publish Thousand Token Wood agent traces to the Hub.

Runs the real multi-model COUNCIL for a few turns, captures each creature's full
prompt, raw response, parsed actions, and private "thought" for every turn --
tagged with which lab's model produced it -- and uploads them as a public
dataset. This is the "Sharing is Caring" bonus quest: open agent traces others
can learn from, and a side-by-side of four labs' small models deciding in the
same market.

Usage (HF_TOKEN must be set, e.g. sourced from .env; engines deployed):
    python scripts/publish_traces.py
    python scripts/publish_traces.py --no-upload     # just write traces.jsonl
"""

import argparse
import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from ttw.actions import _extract_json, parse_actions
from ttw.agents import make_council_policy
from ttw.council import CREATURE_ENGINE, ENGINE_LABELS, ENGINES
from ttw.events import EventDeck
from ttw.llm import build_council_clients
from ttw.sim import step
from ttw.world import seed_world

DATASET_REPO = "build-small-hackathon/thousand-token-wood-traces"
OUT = Path(__file__).resolve().parents[1] / "traces.jsonl"

DATASET_CARD = """\
---
license: mit
language: [en]
tags: [agent-traces, multi-agent, small-models, gradio, build-small-hackathon,
       minicpm, nemotron, gpt-oss]
---

# Thousand Token Wood -- Council Agent Traces

Open agent traces from [Thousand Token Wood](https://huggingface.co/spaces/build-small-hackathon/thousand-token-wood-sim),
a tiny emergent economy where five woodland creatures trade goods for pebbles,
gossip, and react to reskinned market-history "Wood Legends". Each creature
thinks on a **different lab's small model** (distinct-engine budget 29.5B <= 32B):

| Creature | Model | Lab |
|---|---|---|
| Oona (owl) | openai/gpt-oss-20b | OpenAI |
| Bramble (squirrel) | openbmb/MiniCPM3-4B | OpenBMB |
| Fenn (fox) | nvidia/Nemotron-Mini-4B-Instruct | NVIDIA |
| Mossback (tortoise) + Pip (mouse) | AdmiralTaco/ttw-trader-0.5b | fine-tuned (ours) |

Each row is one creature's turn: the exact system + user prompt it saw, its raw
JSON response, its parsed offers and gossip, its private `thought`, and the
model/lab that produced it -- so you can compare how four labs' small models
read the *same* market state and choose differently. Shared for the **Sharing
is Caring** bonus quest of the Build Small Hackathon.

## Fields
- `turn` (int), `creature` (str)
- `engine` (str), `model` (str), `lab` (str): which council engine produced the row
- `system`, `user` (str): the exact prompt the agent received
- `response` (str): the raw model output (normalized; harmony channels stripped)
- `thought` (str|null): the agent's private reasoning for the turn
- `offers` (list): parsed `{creature, side, good, price, qty}` actions
- `gossip` (list[str]): rumors the agent broadcast
"""


def _warm_engines(clients) -> None:
    """Ping every engine once so none cold-starts mid-run (a cold engine can
    hard-timeout and its creatures would sit the turn out)."""
    ping = [[{"role": "user", "content": "ready?"}]]
    for eid, client in clients.items():
        try:
            client.chat_batch(ping, max_tokens=8)
            print(f"  warmed {eid}")
        except Exception as e:  # noqa: BLE001
            print(f"  warm {eid} failed (continuing): {e!r}")


def generate(turns: int = 8) -> list[dict]:
    world = seed_world()
    clients = build_council_clients()
    _warm_engines(clients)
    policy = make_council_policy(clients, temperature=0.7)
    deck = EventDeck(seed=5)
    records: list[dict] = []
    for _ in range(turns):
        if world.turn in (3, 6):  # draw a Wood Legend mid-run for richer traces
            deck.draw(world)
        step(world, policy)
        st = policy.state
        for name, raw in st["raw"].items():
            msgs = st["messages"][name]
            parsed = _extract_json(raw) or {}
            offers, gossip = parse_actions(name, raw)
            eid = CREATURE_ENGINE.get(name, "")
            model_label, lab = ENGINE_LABELS.get(eid, ("", ""))
            records.append(
                {
                    "turn": world.turn,
                    "creature": name,
                    "engine": eid,
                    "model": ENGINES[eid].model if eid in ENGINES else model_label,
                    "lab": lab,
                    "system": msgs[0]["content"],
                    "user": msgs[1]["content"],
                    "response": raw,
                    "thought": parsed.get("thought"),
                    "offers": [vars(o) for o in offers],
                    "gossip": [g.message for g in gossip],
                }
            )
    return records


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--turns", type=int, default=8)
    ap.add_argument("--no-upload", action="store_true")
    ap.add_argument("--force", action="store_true",
                    help="publish even if the engine-health gate trips")
    args = ap.parse_args()

    records = generate(args.turns)
    OUT.write_text(
        "\n".join(json.dumps(r, ensure_ascii=False) for r in records), encoding="utf-8"
    )
    # Engine-health gate: empties per lab, not rows per lab (row counts are a
    # constant of the cast and can never reveal a flaked engine). Scattered
    # empties are expected (e.g. gpt-oss truncation); a whole lab silent or a
    # large empty fraction means the dataset's "four labs side-by-side" claim
    # would be false -- refuse to publish that.
    total = {r["lab"]: 0 for r in records}
    empty = dict.fromkeys(total, 0)
    for r in records:
        total[r["lab"]] += 1
        if not r["response"].strip():
            empty[r["lab"]] += 1
    n_empty = sum(empty.values())
    print(f"wrote {len(records)} trace records to {OUT}")
    print("  empties by lab: " + ", ".join(
        f"{lab}: {empty[lab]}/{total[lab]}" for lab in total))
    dead_labs = [lab for lab in total if empty[lab] == total[lab]]
    too_empty = records and n_empty / len(records) > 0.25
    if (dead_labs or too_empty) and not args.force:
        print(f"ABORTING upload: dead labs {dead_labs or 'none'}, "
              f"empty fraction {n_empty}/{len(records)}. "
              "Fix the engines (or rerun with --force to publish anyway).")
        sys.exit(1)

    if args.no_upload:
        return
    token = os.environ.get("HF_TOKEN")
    if not token:
        print("HF_TOKEN not set; skipping upload. (source .env first)")
        return

    from huggingface_hub import HfApi

    api = HfApi(token=token)
    api.create_repo(DATASET_REPO, repo_type="dataset", exist_ok=True)
    api.upload_file(
        path_or_fileobj=str(OUT),
        path_in_repo="traces.jsonl",
        repo_id=DATASET_REPO,
        repo_type="dataset",
    )
    api.upload_file(
        path_or_fileobj=DATASET_CARD.encode("utf-8"),
        path_in_repo="README.md",
        repo_id=DATASET_REPO,
        repo_type="dataset",
    )
    print(f"published: https://huggingface.co/datasets/{DATASET_REPO}")


if __name__ == "__main__":
    main()