#!/usr/bin/env python3
"""Executable spec for postmortem.analyze_workflows.

Run directly: `python3 postmortem/tests/test_workflow_aggregation.py`.
Exits 0 on pass, 1 on failure. No pytest dependency — same plain-assert
convention as test_terminal_status.py.

Pins issue #28: sessions that ran Workflow-substrate pipelines
(workflows/wf_*.json + subagents/workflows/<runId>/agent-*.jsonl) must
not read as zero subagent spend. The raw transcripts are authoritative
(dedup by requestId); the summary supplies stage labels; a transcript
missing from workflowProgress groups under '?' instead of being folded
into a named stage or dropped.
"""

import json
import os
import sys
import tempfile

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "scripts"))
from postmortem import analyze_workflows  # noqa: E402


def _write(path, obj_lines):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w") as fh:
        for obj in obj_lines:
            fh.write(json.dumps(obj) + "\n")


def _assistant_line(request_id, in_tok, out_tok, model="claude-sonnet-5"):
    return {
        "type": "assistant",
        "requestId": request_id,
        "message": {
            "model": model,
            "usage": {"input_tokens": in_tok, "output_tokens": out_tok,
                      "cache_creation_input_tokens": 0,
                      "cache_read_input_tokens": 0},
        },
    }


def main():
    with tempfile.TemporaryDirectory() as session_dir:
        run_id = "wf_test-1"
        wf = {
            "runId": run_id,
            "workflowName": "spec-fixture",
            "status": "completed",
            "totalTokens": 999,  # live counter: deliberately wrong
            "defaultModel": "claude-sonnet-5",
            "workflowProgress": [
                {"type": "workflow_phase", "index": 1, "title": "P"},
                {"type": "workflow_agent", "agentId": "a1",
                 "label": "impl:1", "model": "claude-sonnet-5",
                 "state": "done", "tokens": 1},
                {"type": "workflow_agent", "agentId": "a2",
                 "label": "impl:2", "model": "claude-sonnet-5",
                 "state": "done", "tokens": 1},
                {"type": "workflow_agent", "agentId": "a4",
                 "label": "qual:1", "model": "claude-opus-4-8",
                 "state": "error", "tokens": 0},  # no transcript on disk
            ],
        }
        os.makedirs(os.path.join(session_dir, "workflows"))
        with open(os.path.join(session_dir, "workflows",
                               run_id + ".json"), "w") as fh:
            json.dump(wf, fh)

        raw = os.path.join(session_dir, "subagents", "workflows", run_id)
        # a1: streamed repeat of the SAME requestId — must dedup to
        # max-per-field (input 10, output 7), not sum to 20/12.
        _write(os.path.join(raw, "agent-a1.jsonl"),
               [_assistant_line("req-1", 10, 5),
                _assistant_line("req-1", 10, 7)])
        # a2: two distinct requests — must sum.
        _write(os.path.join(raw, "agent-a2.jsonl"),
               [_assistant_line("req-2", 3, 2),
                _assistant_line("req-3", 4, 1)])
        # a3: transcript with NO workflowProgress entry — counts under '?'.
        _write(os.path.join(raw, "agent-a3.jsonl"),
               [_assistant_line("req-4", 100, 50)])

        runs = analyze_workflows(session_dir)

    assert len(runs) == 1, runs
    run = runs[0]
    assert run["run_id"] == run_id
    assert run["name"] == "spec-fixture"
    assert run["reported_total_tokens"] == 999

    # Three transcripts on disk, one untracked.
    assert run["agents"] == 3, run
    assert run["agents_untracked"] == 1, run

    stages = run["stages"]
    assert set(stages) == {"impl", "?"}, stages

    impl = stages["impl"]
    assert impl["count"] == 2, impl
    # a1 deduped (10 in / 7 out) + a2 summed (7 in / 3 out)
    assert impl["tokens"]["input"] == 17, impl
    assert impl["tokens"]["output"] == 10, impl
    assert impl["total_tokens"] == 27, impl
    assert impl["models"] == {"claude-sonnet-5": 3}, impl

    unk = stages["?"]
    assert unk["count"] == 1, unk
    assert unk["tokens"]["input"] == 100 and unk["tokens"]["output"] == 50, unk

    # Run totals = authoritative raw aggregation, not the live counter.
    assert run["tokens"]["input"] == 117, run
    assert run["tokens"]["output"] == 60, run
    assert run["total_tokens"] == 177, run

    print("ok: analyze_workflows aggregates raw transcripts per stage, "
          "dedups by requestId, and isolates untracked agents under '?'")
    return 0


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