#!/usr/bin/env python3
"""Muse Wonder Full 1.0 (max), Text-to-Scene re-run: per-case cost from the muse CLI's own session logs.

The harness stream (muse_r1_raw.jsonl) has no usage, but the CLI writes one: every model call ends in a
`model_completed` event with the provider's token counts, in
<attempt>/episode.agent/home/.local/share/muse/sessions/<y>/<m>/<d>/<session>/session.jsonl and, for the
subagents it starts, .../<session>/subagent/<id>/session.jsonl. The kept attempt of each case is
<run_dir>/cases/<case>/final, as the internal cost table uses for every other model.

Semantics, checked on the logs: input_tokens includes cached_tokens (= cache_read_tokens; cache writes are 0)
and output_tokens includes reasoning_tokens. Prices per million tokens, OpenRouter's catalogue (the source the
internal table uses for token-priced models): meta/muse-wonder $4 input, $0.20 cached input, $20 output (Wonder
Full is not sold per token; Muse Wonder's rate is the nearest list price), meta/muse-spark-1.3 $1.25 / $0.15 /
$4.25 for the CLI's background model. The CLI's own cost ledger (its session view snapshot) gives the same totals.
Not included: the skill-reminder side agent's decision calls (about one per main step), whose usage the CLI does
not write to any file the run keeps readable. Read-only; JSON on stdout. The run folders are the arguments:
    ssh <host> python3 - <run_dir> ... < scripts/t2s_muse_cost.py > ../t2s-src/musewf_cost.json
"""
import glob
import json
import os
import statistics
import sys

RUNS = sys.argv[1:]
PRICE = {"muse-wonder-full-1.0": (4.0, 0.20, 20.0, "openrouter meta/muse-wonder"),
         "muse-spark-1.3": (1.25, 0.15, 4.25, "openrouter meta/muse-spark-1.3")}


def calls(path):
    for line in open(path, errors="ignore"):
        if '"model_completed"' not in line:
            continue
        try:
            ev = (json.loads(line).get("payload") or {}).get("event") or {}
        except ValueError:
            continue
        if isinstance(ev, dict) and ev.get("kind") == "model_completed":
            yield ev.get("model"), ev.get("usage") or {}


out, problems = {}, []
for rd in RUNS:
    for cdir in sorted(glob.glob(os.path.join(rd, "cases", "*-long"))):
        case = os.path.basename(cdir)
        att = os.path.realpath(os.path.join(cdir, "final"))
        if not os.path.isdir(os.path.join(att, "episode")):
            continue
        logs = sorted(glob.glob(os.path.join(att, "episode.agent/home/.local/share/muse/sessions/2*/*/*/*/session.jsonl")))
        subs = sorted(glob.glob(os.path.join(att, "episode.agent/home/.local/share/muse/sessions/2*/*/*/*/subagent/*/session.jsonl")))
        tot, n, usd = {}, {}, 0.0
        for path in logs + subs:
            for model, u in calls(path):
                if model not in PRICE:
                    problems.append("%s: unpriced model %s" % (case, model))
                    continue
                inp, cached, outp = (int(u.get(k) or 0) for k in ("input_tokens", "cached_tokens", "output_tokens"))
                if cached > inp or (u.get("reasoning_tokens") or 0) > outp or (u.get("cache_write_tokens") or 0):
                    problems.append("%s: usage outside the checked semantics %s" % (case, u))
                p_in, p_cached, p_out = PRICE[model][:3]
                usd += ((inp - cached) * p_in + cached * p_cached + outp * p_out) / 1e6
                t = tot.setdefault(model, {"input": 0, "cached": 0, "output": 0, "reasoning": 0})
                t["input"] += inp; t["cached"] += cached; t["output"] += outp; t["reasoning"] += int(u.get("reasoning_tokens") or 0)
                n[model] = n.get(model, 0) + 1
        if not logs:
            problems.append("%s: no CLI session log in %s" % (case, att))
        out[case] = {"attempt": att, "sessions": len(logs), "subagent_sessions": len(subs), "calls": n, "tokens": tot,
                     "cost_usd": round(usd, 6) if logs else None}
costs = [v["cost_usd"] for v in out.values() if v["cost_usd"] is not None]
json.dump({"method": __doc__, "prices_per_million": PRICE, "cases": out, "problems": problems,
           "median_cost_usd": statistics.median(costs) if costs else None, "n": len(costs)}, sys.stdout, indent=1)