"""Test SwarmRunMetricsCollector against a synthetic run with known states.

Hermetic (REDIS_FAKE). Run from agent_swarm_v6:  python scripts/test-benchmark-collector.py
"""
import asyncio
import json
import math
import os
import sys
from pathlib import Path

os.environ["REDIS_FAKE"] = "1"
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from orchestrator.redis_client import redis_client
from orchestrator import swarm_runtime as sr_mod
from orchestrator.swarm_runtime import swarm_runtime
from orchestrator.task_queue import task_queue, TaskStatus
from benchmark.collectors.run_collector import SwarmRunMetricsCollector

failures = []


def check(name, cond):
    print(("PASS" if cond else "FAIL"), "-", name)
    if not cond:
        failures.append(name)


async def _noop(self, *a, **k):
    return None


async def add_task(run, task_id, *, status, agent, cost=0.0, retry=0, depends_on=None):
    t = await task_queue.create_task(task_id=task_id, description=task_id, agent_role=task_id.split("-")[-1],
                                     depends_on=depends_on or [], enqueue=False)
    t.status = status
    t.assigned_agent_id = agent
    t.retry_count = retry
    t.result = json.dumps({"usage": {"model_cost_usd": cost}})
    await task_queue._save_task(t)
    await swarm_runtime.attach_task(run, t.task_id)
    return t


async def main():
    await redis_client.connect()
    sr_mod.SwarmRuntime._post_callback = _noop  # no real HTTP

    body = {
        "mode": "swarm",
        "requirement": {"objective": "collector test"},
        "orchestration_plan": {"budget": {"max_cost_usd": 10}},
        "callback": {"url": "http://127.0.0.1:9/cb", "subscribed_events": []},
        "metadata": {"manager_deployment_id": "m-col"},
    }
    run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="c")

    await add_task(run, "t-implementation", status=TaskStatus.COMPLETED, agent="A", cost=2.0)
    await add_task(run, "t-testing", status=TaskStatus.COMPLETED, agent="B", cost=3.0, retry=1,
                   depends_on=["t-implementation"])
    await add_task(run, "t-documentation", status=TaskStatus.FAILED, agent="A", cost=0.0, retry=3)

    await swarm_runtime.emit_event(run, "handoff.requested", payload={"task_id": "t1", "from_role": "implementation", "to_role": "testing"})
    await swarm_runtime.emit_event(run, "handoff.completed", payload={"task_id": "t1", "from_role": "implementation", "to_role": "testing"})

    # Peer communication: 3 requests routed, 2 of them answered (matching correlation_ids).
    for cid, delivered in [("c1", True), ("c2", True), ("c3", True)]:
        await swarm_runtime.record_peer_message(run, correlation_id=cid, is_reply=False, delivered=delivered)
    for cid in ("c1", "c2"):  # c3 never gets a reply
        await swarm_runtime.record_peer_message(run, correlation_id=cid, is_reply=True, delivered=True)

    collector = SwarmRunMetricsCollector(run.swarm_id)
    metrics = await collector.collect()
    cov = collector.coverage

    # s_completion = 2/3*100
    check("s_completion = 66.67", round(metrics.s_completion, 2) == 66.67 and cov["s_completion"])
    # s_collaboration = 0.5*100(handoff) + 0.3*100(dep resolved) + 0.2*50(workload A:2,B:1) = 90
    check("s_collaboration = 90.0", round(metrics.s_collaboration, 1) == 90.0 and cov["s_collaboration"])
    # s_robustness: failures={t-testing(retry),t-documentation(failed)}=2, recovered={t-testing completed}=1 => 50
    check("s_robustness = 50.0", round(metrics.s_robustness, 1) == 50.0 and cov["s_robustness"])
    # s_cost = 100*budget(10)/actual(5) = 200
    check("s_cost = 200.0", round(metrics.s_cost, 1) == 200.0 and cov["s_cost"])
    # governance: no approvals -> NaN, coverage False
    check("s_governance NaN + coverage False", math.isnan(metrics.s_governance) and cov["s_governance"] is False)
    # s_communication = 2 answered / 3 requests * 100 = 66.67
    check("s_communication = 66.67", round(metrics.s_communication, 2) == 66.67 and cov["s_communication"])
    # not-yet-collectable -> NaN + coverage False
    check("uncollectable metrics NaN + coverage False",
          all(math.isnan(getattr(metrics, k)) and cov[k] is False
              for k in ("tau", "eta", "p_decision", "reward", "s_gain",
                        "s_swarm", "g_e", "g_e_cost", "benchmark")))

    # --- governance telemetry counter: REAL when the run had governed ops (#23) ---
    # CompliantOperations/TotalOperations: 2 approvals decided (approved+rejected) of 3 governed
    # ops -> 66.67. Proves the governance counter is wired, not just the empty/NaN path above.
    gov_run, _ = await swarm_runtime.get_or_create_run(
        body={**body, "metadata": {"manager_deployment_id": "m-gov"}},
        idempotency_key=None, correlation_id="cg")
    gov_run.approvals = {
        "a1": {"approval_id": "a1", "decision": "approved"},
        "a2": {"approval_id": "a2", "decision": "rejected"},
        "a3": {"approval_id": "a3", "decision": "pending"},
    }
    await add_task(gov_run, "g-implementation", status=TaskStatus.COMPLETED, agent="A", cost=1.0)
    await swarm_runtime.save_run(gov_run)
    gov_collector = SwarmRunMetricsCollector(gov_run.swarm_id)
    gm = await gov_collector.collect()
    check("s_governance = 66.67 (2 decided / 3 governed ops)",
          round(gm.s_governance, 2) == 66.67 and gov_collector.coverage["s_governance"] is True)

    print()
    if failures:
        print(f"{len(failures)} collector check(s) FAILED: {failures}")
        sys.exit(1)
    print("all benchmark collector checks passed")


if __name__ == "__main__":
    asyncio.run(main())