"""Pattern 5: a supervised swarm. Peers investigate competing hypotheses, share a blackboard, and a judge decides."""

import asyncio
import json
import shutil
import tempfile
from dataclasses import dataclass, field
from pathlib import Path

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

# work on a pristine copy so the bugs are present regardless of earlier demos
WORK = Path(tempfile.mkdtemp(prefix="debate-")) / "project"
shutil.copytree(PRISTINE, WORK)


async def run(prompt, schema=None, **opts):
    options = ClaudeAgentOptions(cwd=str(WORK), setting_sources=[], max_turns=opts.pop("max_turns", 6),
                                 output_format={"type": "json_schema", "schema": schema} if schema else None, **opts)
    async for m in query(prompt=prompt, options=options):
        if isinstance(m, ResultMessage):
            return (m.structured_output if schema else m.result), (m.total_cost_usd or 0)


@dataclass
class Blackboard:
    hypotheses: dict[str, str]
    evidence: dict[str, list[str]] = field(default_factory=dict)
    cost: float = 0.0


EVIDENCE_SCHEMA = {"type": "object",
                   "properties": {"for": {"type": "array", "items": {"type": "string"}},
                                  "against": {"type": "object", "additionalProperties": {"type": "array", "items": {"type": "string"}}}},
                   "required": ["for", "against"]}
VERDICT_SCHEMA = {"type": "object", "properties": {"winner": {"type": "string"}, "reason": {"type": "string"},
                                                   "confidence": {"type": "number"}}, "required": ["winner", "reason", "confidence"]}


async def investigate(bb: Blackboard, name: str, theory: str):
    others = {n: t for n, t in bb.hypotheses.items() if n != name}
    out, c = await run(
        f"You are investigator {name}. The failing test is tests/test_utils.py::test_user_name_none_returns_empty "
        f"(it calls get_user_name(None) and expects ''). Your theory: {theory}.\nOther open theories: {json.dumps(others)}\n"
        f"Evidence already on the blackboard:\n{json.dumps(bb.evidence, indent=2)}\n"
        "Read the code and the test. Return evidence FOR your theory and evidence AGAINST each other theory "
        "(cite file:line). Be honest: if your theory is wrong, say so in 'for' with an empty list.",
        EVIDENCE_SCHEMA, allowed_tools=["Read", "Grep", "Glob"], effort="low")
    bb.cost += c
    bb.evidence.setdefault(name, []).extend(out["for"])
    for other, reasons in out["against"].items():
        bb.evidence.setdefault(f"against_{other}", []).extend(reasons)


async def debate(bb: Blackboard, rounds: int = 2):
    for r in range(rounds):                                                             # (1) bounded rounds
        await asyncio.gather(*(investigate(bb, n, t) for n, t in bb.hypotheses.items()))  # (2) peers in parallel, shared blackboard
        print(f"after round {r + 1}: evidence counts = { {k: len(v) for k, v in bb.evidence.items()} }")
    verdict, c = await run("Judge which hypothesis best explains the failing test given this evidence. "
                           "Prefer the theory with concrete file:line support and no rebuttals.\n" + json.dumps(bb.evidence, indent=2),
                           VERDICT_SCHEMA, tools=[])                                   # (3) independent judge, no tools
    bb.cost += c
    return verdict


async def main():
    bb = Blackboard(hypotheses={
        "A": "get_user_name crashes because it subscripts user['name'] without checking for None",
        "B": "the test fails because of an import error: src.utils cannot be imported",
        "C": "calculate_average's division by zero is what makes the whole test module fail",
    })
    verdict = await debate(bb)
    print("\nBLACKBOARD:")
    for k, v in bb.evidence.items():
        print(f"  {k}:")
        for item in v[:3]:
            print(f"    - {item[:140]}")
    print(f"\nVERDICT: {verdict}")
    print(f"TOTAL COST: ${bb.cost:.4f}")


asyncio.run(main())