#!/usr/bin/env python3
"""mkfixture.py -- synthetic TRACE_FORMAT.md logs for testing tracecmp.py.

usage:  mkfixture.py OUTDIR [--seed N] [--calls N]

Writes to OUTDIR:
  trace_clean.jsonl        golden `trace` log (meta + N calls over 3 hooks)
  compare_clean.jsonl      `compare` log, ours == original everywhere
  compare_bad.jsonl        `compare` log with injected divergences (see INJECTED)
  replay_impl_ok.jsonl     impl output matching trace_clean (for --replay)
  replay_impl_bad.jsonl    impl output with one wrong ret and one missing call
  invalid.jsonl            broken lines: bad JSON, missing field, bad type, dup id

The records are written by emit_record(), a Python transliteration of the
C++ emitter rules in TRACE_FORMAT.md section 4 (fixed key order, hand-rolled
string escaping, %.9g floats) -- NOT json.dumps -- so the tests prove that an
emitter following those rules produces logs the harness reads back exactly.

Stdlib only.
"""
from __future__ import annotations

import argparse
import copy
import hashlib
import os
import random
import struct
import sys
from typing import Any

KEY_ORDER = ["ts", "hook", "mode", "call_id", "thread", "depth", "args", "ret", "side",
             "ours", "diverged", "diff", "err", "note"]
INLINE_MAX = 256
HOOKS = ["CfgVar_RegisterKey", "Manifest_Load", "Mars::ParseBlock"]


# --- typed-value constructors -----------------------------------------------------

def i32(v, n=None):
    return _tv("i32", v, n)


def u32(v, n=None):
    return _tv("u32", v, n)


def u64(v, n=None):
    return _tv("u64", v if v < 2**53 else str(v), n)


def f32(v, n=None):
    if v != v:
        return _tv("f32", "nan", n)
    if v in (float("inf"), float("-inf")):
        return _tv("f32", "inf" if v > 0 else "-inf", n)
    return _tv("f32", round_f32(v), n)


def boolean(v, n=None):
    return _tv("bool", bool(v), n)


def s(v: str, n=None):
    """cp1252 byte string -> str tv (each byte one code point)."""
    return _tv("str", v.encode("cp1252").decode("latin-1"), n)


def ptr(addr: int, n=None):
    return _tv("ptr", "0x%08x" % addr, n)


def enum(v: int, name: str | None = None, n=None):
    tv = _tv("enum", v, n)
    if name:
        tv["name"] = name
    return tv


def null():
    return {"t": "null", "v": None}


def lst(items, n=None):
    return _tv("list", list(items), n)


def sset(items, n=None):
    return _tv("set", list(items), n)


def struct_(fields: dict, n=None):
    return _tv("struct", dict(fields), n)


def jsonv(v, n=None):
    return _tv("json", v, n)


def by(data: bytes, inline_max: int = INLINE_MAX):
    """bytes tv: n + sha256 always; hex inline when small, else head.  (No 'n'
    name on bytes -- the key is taken by the length; name the side region instead.)"""
    tv: dict = {"t": "bytes", "n": len(data), "sha256": hashlib.sha256(data).hexdigest()}
    if len(data) <= inline_max:
        tv["hex"] = data.hex()
    else:
        tv["head"] = data[:32].hex()
    return tv


def _tv(t, v, n):
    tv = {"t": t, "v": v}
    if n:
        tv["n"] = n
    return tv


def round_f32(x: float) -> float:
    return struct.unpack("<f", struct.pack("<f", x))[0]


# --- reference emitter (mirrors the C++ rules) --------------------------------------

def esc(u: str, wide: bool = False) -> str:
    out = ['"']
    for c in u:
        o = ord(c)
        if c == '"':
            out.append('\\"')
        elif c == "\\":
            out.append("\\\\")
        elif c == "\n":
            out.append("\\n")
        elif c == "\r":
            out.append("\\r")
        elif c == "\t":
            out.append("\\t")
        elif o < 0x20 or o >= 0x7F:
            out.append("\\u%04x" % o)
        else:
            out.append(c)
    out.append('"')
    return "".join(out)


def fmt_num(v: Any, t: str | None = None) -> str:
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, int):
        return "%d" % v
    if isinstance(v, float):
        if v != v:
            return '"nan"'
        if v in (float("inf"), float("-inf")):
            return '"inf"' if v > 0 else '"-inf"'
        return ("%.9g" if t == "f32" else "%.17g") % v
    raise TypeError(v)


def emit_value(x: Any, t: str | None = None, plain: bool = False) -> str:
    """Generic JSON value emitter with the spec's escaping.  Nested dicts that
    look like typed values are emitted as such unless plain=True (json payloads)."""
    if x is None:
        return "null"
    if isinstance(x, (bool, int, float)):
        return fmt_num(x, t)
    if isinstance(x, str):
        return esc(x)
    if isinstance(x, list):
        return "[" + ",".join(emit_value(e, plain=plain) for e in x) + "]"
    if isinstance(x, dict):
        if not plain and "t" in x and ("v" in x or x["t"] == "bytes"):
            return emit_tv(x)
        return "{" + ",".join(esc(k) + ":" + emit_value(v, plain=plain) for k, v in x.items()) + "}"
    raise TypeError(x)


def emit_tv(tv: dict) -> str:
    t = tv["t"]
    parts = ['"t":' + esc(t)]
    if t == "bytes":
        parts.append('"n":%d' % tv["n"])
        parts.append('"sha256":' + esc(tv["sha256"]))
        if "hex" in tv:
            parts.append('"hex":' + esc(tv["hex"]))
        if "head" in tv:
            parts.append('"head":' + esc(tv["head"]))
    else:
        v = tv["v"]
        if t in ("list", "set"):
            parts.append('"v":[' + ",".join(emit_tv(e) for e in v) + "]")
        elif t == "struct":
            parts.append('"v":{' + ",".join(esc(k) + ":" + emit_tv(e) for k, e in v.items()) + "}")
        elif t == "json":
            parts.append('"v":' + emit_value(v, plain=True))
        elif t in ("f32", "f64"):
            parts.append('"v":' + (esc(v) if isinstance(v, str) else fmt_num(v, t)))
        elif t == "wstr":
            parts.append('"v":' + esc(v, wide=True))
        elif v is None:
            parts.append('"v":null')
        elif isinstance(v, str):
            parts.append('"v":' + esc(v))
        else:
            parts.append('"v":' + fmt_num(v))
        if "name" in tv:
            parts.append('"name":' + esc(tv["name"]))
    if "n" in tv and t != "bytes":
        parts.append('"n":' + esc(tv["n"]))
    return "{" + ",".join(parts) + "}"


def emit_side(side: dict) -> str:
    ents = []
    for name, ent in side.items():
        inner = []
        if "before" in ent:
            inner.append('"before":' + ("null" if ent["before"] is None else emit_tv(ent["before"])))
        inner.append('"after":' + emit_tv(ent["after"]))
        ents.append(esc(name) + ":{" + ",".join(inner) + "}")
    return "{" + ",".join(ents) + "}"


def emit_record(rec: dict) -> str:
    parts = []
    for k in KEY_ORDER:
        if k not in rec:
            continue
        v = rec[k]
        if k in ("ts", "call_id", "thread", "depth"):
            parts.append(esc(k) + ":" + fmt_num(v))
        elif k in ("hook", "mode", "err", "note"):
            parts.append(esc(k) + ":" + esc(v))
        elif k == "args":
            parts.append('"args":[' + ",".join(emit_tv(a) for a in v) + "]")
        elif k == "ret":
            parts.append('"ret":' + ("null" if v is None else emit_tv(v)))
        elif k == "side":
            parts.append('"side":' + emit_side(v))
        elif k == "ours":
            parts.append('"ours":{"ret":' + ("null" if v["ret"] is None else emit_tv(v["ret"])) +
                         ',"side":' + emit_side(v["side"]) + "}")
        elif k == "diverged":
            parts.append('"diverged":' + ("true" if v else "false"))
        elif k == "diff":
            parts.append('"diff":' + emit_value(v))
        else:
            raise AssertionError(k)
    return "{" + ",".join(parts) + "}\n"


def emit_meta(meta: dict) -> str:
    return '{"meta":' + emit_value(meta) + "}\n"


def write_log(path: str, meta: dict | None, records: list[dict], raw_lines: list[str] | None = None) -> None:
    with open(path, "w", encoding="utf-8", newline="\n") as f:
        if meta is not None:
            f.write(emit_meta(meta))
        for r in records:
            f.write(emit_record(r))
        for line in raw_lines or []:
            f.write(line if line.endswith("\n") else line + "\n")


# --- scenario generation ----------------------------------------------------------------

def meta(**hooks) -> dict:
    return {"format": 1, "build": "mkfixture 1", "exe_sha256": "0" * 64,
            "started": "2026-09-07T00:00:00Z", "inline_max": INLINE_MAX, "hooks": hooks}


def gen_calls(n: int, seed: int = 1) -> list[dict]:
    """n trace calls cycling over HOOKS, with every tv type represented."""
    rng = random.Random(seed)
    recs = []
    for cid in range(n):
        hook = HOOKS[cid % len(HOOKS)]
        ts = 1000 + cid * 37
        if hook == "CfgVar_RegisterKey":
            key = rng.choice(["ForceSingleCore", "Résumé", "Tab\tKey", 'Quote"d', "Back\\slash"])
            blob = bytes(rng.getrandbits(8) for _ in range(rng.choice([0, 16, 512])))
            rec = {"ts": ts, "hook": hook, "mode": "trace", "call_id": cid, "thread": 4120, "depth": 0,
                   "args": [s(key, "key"), i32(rng.randint(-5, 5), "value"), ptr(0x00a36fd0 + cid * 16, "table")],
                   "ret": boolean(True),
                   "side": {"cfg_table": {"before": by(blob), "after": by(blob[::-1])}}}
        elif hook == "Manifest_Load":
            ids = [rng.randint(1, 400) for _ in range(rng.randint(1, 5))]
            rec = {"ts": ts, "hook": hook, "mode": "trace", "call_id": cid, "thread": 4120,
                   "args": [s("Weapons/_weapons.txt", "path")],
                   "ret": i32(len(ids), "count"),
                   "side": {"registry": {"after": struct_({
                       "entries": lst(struct_({"id": u32(i), "name": s("w%d.weapon" % i)}) for i in ids),
                       "deleted": sset(u32(i) for i in rng.sample(range(500, 600), 3)),
                       "flags": enum(rng.randint(0, 3), "MF_%d" % 0),
                       "big": u64(2**60 + cid),
                       "nothing": null()})}}}
        else:
            vals = [round_f32(rng.uniform(-10, 10)) for _ in range(3)]
            rec = {"ts": ts, "hook": hook, "mode": "trace", "call_id": cid, "thread": 4124, "depth": 1,
                   "args": [s("weapon { damage %g }" % vals[0], "text"), u32(cid, "len")],
                   "ret": jsonv({"weapon": {"damage": vals[0], "range": vals[1], "name": "Laseré", "count": 3,
                                            "ok": True, "tags": ["a", "b"]}}),
                   "side": {"scratch": {"after": lst(f32(v) for v in vals)},
                            "nanbox": {"after": f32(float("nan"))}}}
        recs.append(rec)
    return recs


def to_compare(recs: list[dict]) -> list[dict]:
    """Clone trace records into compare records where ours == original."""
    out = []
    for r in recs:
        c = dict(r)
        c["mode"] = "compare"
        c["ours"] = {"ret": _deep(r["ret"]), "side": {k: {"after": _deep(v["after"])} for k, v in r["side"].items()}}
        c["diverged"] = False
        c["diff"] = []
        out.append(c)
    return out


def _deep(x):
    return copy.deepcopy(x)


def inject_divergences(recs: list[dict]) -> tuple[list[dict], list[dict]]:
    """Mutate `ours` on a handful of compare records (the INJECTED cases below).
    Returns (records, expectations); expectation = {call_id, diverges_exact,
    diverges_tol (abs 1e-3), why, path}.  Needs >= 12 records (4 per hook)."""
    recs = [_deep(r) for r in recs]
    exp = []
    by_hook: dict[str, list[dict]] = {}
    for r in recs:
        by_hook.setdefault(r["hook"], []).append(r)

    cfg = by_hook["CfgVar_RegisterKey"]
    # 1. bool ret flipped
    cfg[0]["ours"]["ret"] = boolean(False)
    exp.append({"call_id": cfg[0]["call_id"], "diverges_exact": True, "diverges_tol": True, "why": "exact", "path": "ret"})
    # 2. bytes region differs (same length, different content) -> hash
    orig = cfg[1]["side"]["cfg_table"]["after"]
    n = orig["n"]
    if n == 0:  # make the region non-empty on both sides so a content diff is possible
        n = 16
        cfg[1]["side"]["cfg_table"]["after"] = orig = by(bytes(range(n)))
    data = bytes(n)
    if data.hex() == orig.get("hex"):
        data = bytes([255]) + data[1:]
    cfg[1]["ours"]["side"]["cfg_table"]["after"] = by(data)
    exp.append({"call_id": cfg[1]["call_id"], "diverges_exact": True, "diverges_tol": True, "why": "hash",
                "path": "side.cfg_table.after"})
    # 3. shim lies: says diverged=true on an identical record -> harness warning only
    cfg[2]["diverged"] = True
    exp.append({"call_id": cfg[2]["call_id"], "diverges_exact": False, "diverges_tol": False, "why": None,
                "path": None, "warn": "shim said diverged"})

    man = by_hook["Manifest_Load"]
    # 4. set reordered -> NOT a divergence
    st = man[0]["ours"]["side"]["registry"]["after"]
    st["v"]["deleted"]["v"] = list(reversed(st["v"]["deleted"]["v"]))
    exp.append({"call_id": man[0]["call_id"], "diverges_exact": False, "diverges_tol": False, "why": None, "path": None})
    # 5. list reordered (2+ entries) or entry id changed -> divergence
    st = man[1]["ours"]["side"]["registry"]["after"]
    ent = st["v"]["entries"]["v"]
    if len(ent) >= 2 and ent[0] != ent[-1]:
        ent.reverse()
    else:
        ent[0]["v"]["id"] = u32(ent[0]["v"]["id"]["v"] + 1)
    exp.append({"call_id": man[1]["call_id"], "diverges_exact": True, "diverges_tol": True, "why": "exact",
                "path": "side.registry.after.v.entries"})
    # 6. struct field missing
    del man[2]["ours"]["side"]["registry"]["after"]["v"]["big"]
    exp.append({"call_id": man[2]["call_id"], "diverges_exact": True, "diverges_tol": True, "why": "missing",
                "path": "side.registry.after.v.big"})

    mp = by_hook["Mars::ParseBlock"]
    # 7. f32 nudged by 1e-5 -> diverges exact, passes abs 1e-3
    sc = mp[0]["ours"]["side"]["scratch"]["after"]["v"]
    sc[1] = f32(sc[1]["v"] + 1e-5)
    exp.append({"call_id": mp[0]["call_id"], "diverges_exact": True, "diverges_tol": False, "why": "exact",
                "path": "side.scratch.after.v[1]"})
    # 8. json float nudged by 0.5 -> diverges both
    rv = mp[1]["ours"]["ret"] = _deep(mp[1]["ret"])
    rv["v"]["weapon"]["range"] = rv["v"]["weapon"]["range"] + 0.5
    exp.append({"call_id": mp[1]["call_id"], "diverges_exact": True, "diverges_tol": True, "why": "exact",
                "path": "ret.weapon.range"})
    # 9. err record
    mp[2]["err"] = "ours threw std::bad_alloc"
    del mp[2]["ours"]
    exp.append({"call_id": mp[2]["call_id"], "diverges_exact": True, "diverges_tol": True, "why": "err", "path": "call"})
    # 10. ptr differs -> NOT a divergence (default policy); type mismatch i32 vs u32 -> divergence
    if len(cfg) > 3:
        cfg[3]["ours"]["ret"] = boolean(True)
        cfg[3]["ours"]["side"]["cfg_table"]["after"] = cfg[3]["side"]["cfg_table"]["after"]
        exp.append({"call_id": cfg[3]["call_id"], "diverges_exact": False, "diverges_tol": False, "why": None, "path": None})
    if len(man) > 3:
        man[3]["ours"]["ret"] = u32(man[3]["ret"]["v"])
        exp.append({"call_id": man[3]["call_id"], "diverges_exact": True, "diverges_tol": True, "why": "type", "path": "ret"})
    # the shim's own verdict is truthful for every injected case except #3
    by_id = {r["call_id"]: r for r in recs}
    for e in exp:
        if "warn" not in e and "diverged" in by_id[e["call_id"]]:
            by_id[e["call_id"]]["diverged"] = e["diverges_exact"]
    return recs, exp


def impl_records(recs: list[dict]) -> list[dict]:
    """Implementation-output records (replay input) equal to the golden outputs."""
    return [{"call_id": r["call_id"], "hook": r["hook"], "ret": r["ret"],
             "side": {k: v["after"] for k, v in r["side"].items()}} for r in recs]


def emit_impl(path: str, recs: list[dict]) -> None:
    with open(path, "w", encoding="utf-8", newline="\n") as f:
        for r in recs:
            side = "{" + ",".join(esc(k) + ":" + emit_tv(v) for k, v in r["side"].items()) + "}"
            f.write('{"call_id":%d,"hook":%s,"ret":%s,"side":%s}\n' % (
                r["call_id"], esc(r["hook"]), "null" if r["ret"] is None else emit_tv(r["ret"]), side))


def invalid_lines() -> list[str]:
    return [
        '{"ts":1,"hook":"X","mode":"trace","call_id":900,"thread":1,"args":[],"ret":null,"side":{}}',  # ok
        '{"ts":1,"hook":"X","mode":"trace","call_id":901,"thread":1,"args":[],"ret":null,"side":{}',   # bad JSON
        '{"ts":1,"hook":"X","mode":"trace","call_id":902,"thread":1,"args":[]}',                        # missing ret/side
        '{"ts":1,"hook":"X","mode":"trace","call_id":903,"thread":1,"args":[{"t":"i32","v":"x"}],"ret":null,"side":{}}',  # bad int
        '{"ts":1,"hook":"X","mode":"trace","call_id":900,"thread":1,"args":[],"ret":null,"side":{}}',  # dup id
        '{"ts":1,"hook":"X","mode":"compare","call_id":904,"thread":1,"args":[],"ret":null,"side":{}}',  # compare w/o ours
        '{"ts":1,"hook":"X","mode":"trace","call_id":905,"thread":1,"args":[{"t":"bytes","n":2,"sha256":"zz"}],"ret":null,"side":{}}',
        '{"ts":1,"hook":"X","mode":"trace","call_id":906,"thread":1,"args":[],"ret":{"t":"u8","v":300},"side":{}}',
        '{"ts":1,"hook":"X","mode":"trace","call_id":907,"thread":1,"args":[],"ret":null,"side":{},"bogus":1}',  # ok + warning
    ]


def build_all(outdir: str, n: int = 12, seed: int = 1) -> dict:
    os.makedirs(outdir, exist_ok=True)
    trace = gen_calls(n, seed)
    m = meta(**{"Mars::ParseBlock": {"ftol": 0}})
    paths = {k: os.path.join(outdir, k + ".jsonl") for k in
             ("trace_clean", "compare_clean", "compare_bad", "replay_impl_ok", "replay_impl_bad", "invalid")}
    write_log(paths["trace_clean"], m, trace)
    clean = to_compare(trace)
    write_log(paths["compare_clean"], m, clean)
    bad, exp = inject_divergences(clean)
    write_log(paths["compare_bad"], m, bad)
    impl = impl_records(trace)
    emit_impl(paths["replay_impl_ok"], impl)
    impl_bad = _deep(impl)
    impl_bad[0]["ret"] = boolean(False)          # wrong ret on call 0
    dropped = impl_bad.pop(1)["call_id"]         # missing call 1
    impl_bad.append({"call_id": 9999, "hook": "Extra", "ret": None, "side": {}})  # extra -> warning
    emit_impl(paths["replay_impl_bad"], impl_bad)
    write_log(paths["invalid"], m, [], invalid_lines())
    return {"paths": paths, "trace": trace, "compare_bad": bad, "expect": exp,
            "replay_bad": {"wrong": 0, "missing": dropped, "extra": 9999}}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("outdir")
    ap.add_argument("--seed", type=int, default=1)
    ap.add_argument("--calls", type=int, default=12)
    a = ap.parse_args(argv)
    r = build_all(a.outdir, a.calls, a.seed)
    for k, p in r["paths"].items():
        print(f"{k:16s} {p}")
    print(f"injected divergences: {sum(1 for e in r['expect'] if e['diverges_exact'])} "
          f"(tolerant: {sum(1 for e in r['expect'] if e['diverges_tol'])})")
    return 0


if __name__ == "__main__":
    sys.exit(main())