"""Pattern 2: a code-owned graph. Your program is the orchestrator; each node is a query().

    plan ──► fan-out (one worker per file, in parallel) ──► synthesize ──► verify
                 ▲                                                          │
                 └──────────── conditional edge: rejected? retry once ◄─────┘

Nodes exchange typed state (dicts validated by JSON Schema), not prose."""

import asyncio
import json
from dataclasses import dataclass, field

from claude_agent_sdk import ClaudeAgentOptions, ResultMessage, query
from demo_util import PROJECT


@dataclass
class State:
    files: list[str] = field(default_factory=list)
    findings: dict[str, list] = field(default_factory=dict)
    report: str = ""
    verdict: dict = field(default_factory=dict)
    cost: float = 0.0
    trace: list[str] = field(default_factory=list)


async def node(state: State, name: str, prompt: str, schema: dict | None = None, **opts):
    """Run one graph node as its own fresh-context agent and return its typed output."""
    options = ClaudeAgentOptions(
        cwd=PROJECT, setting_sources=[], max_turns=opts.pop("max_turns", 6),
        output_format={"type": "json_schema", "schema": schema} if schema else None, **opts,
    )
    out, text = None, None
    async for m in query(prompt=prompt, options=options):
        if isinstance(m, ResultMessage):
            out, text = m.structured_output, m.result
            state.cost += m.total_cost_usd or 0
            state.trace.append(f"{name}: {m.subtype}, {m.num_turns} turns, ${m.total_cost_usd:.4f}")
    return out if schema else text


FILES_SCHEMA = {"type": "object", "properties": {"files": {"type": "array", "items": {"type": "string"}}},
                "required": ["files"]}
FINDINGS_SCHEMA = {"type": "object", "properties": {"findings": {"type": "array", "items": {
    "type": "object", "properties": {"line": {"type": "integer"}, "issue": {"type": "string"},
                                     "severity": {"type": "string", "enum": ["low", "medium", "high"]}},
    "required": ["line", "issue", "severity"]}}}, "required": ["findings"]}
VERDICT_SCHEMA = {"type": "object", "properties": {"approved": {"type": "boolean"}, "reason": {"type": "string"}},
                  "required": ["approved", "reason"]}


async def plan(state: State):
    out = await node(state, "plan", "List every Python source file under src/ (relative paths).",
                     FILES_SCHEMA, allowed_tools=["Glob"])
    state.files = out["files"]


async def fan_out(state: State):
    async def worker(path):
        out = await node(state, f"worker[{path}]",
                         f"Review {path} for bugs and security issues. Report findings with line numbers.",
                         FINDINGS_SCHEMA, allowed_tools=["Read"], effort="low")
        return path, out["findings"]
    results = await asyncio.gather(*(worker(p) for p in state.files))   # parallel edges
    state.findings = dict(results)


async def synthesize(state: State):
    state.report = await node(
        state, "synthesize",
        "Write a short prioritized review (max 8 bullets) from these findings. Do not read files.\n"
        + json.dumps(state.findings, indent=2), tools=[])


async def verify(state: State):
    state.verdict = await node(
        state, "verify",
        "You are an independent reviewer with fresh context. Approve this report only if every finding "
        "cites a line number and the highest-severity items come first. Report:\n" + state.report,
        VERDICT_SCHEMA, tools=[])


async def main():
    state = State()
    await plan(state)
    print("plan ->", state.files)
    await fan_out(state)
    print("fan_out ->", {k: len(v) for k, v in state.findings.items()})
    for attempt in range(2):                               # bounded retry edge
        await synthesize(state)
        await verify(state)
        print(f"verify (attempt {attempt + 1}) ->", state.verdict)
        if state.verdict["approved"]:
            break
    print("\nREPORT:\n" + state.report)
    print("\nTRACE:"); print("\n".join("  " + t for t in state.trace))
    print(f"TOTAL COST: ${state.cost:.4f}")


asyncio.run(main())