"""Usage review-retry cost attribution test (issue #16).

Verifies emit_usage_event tags each usage event with `cost_phase`:
  * "initial" before cross-review has reopened the task for rework;
  * "review_retry" once the task appears in run.metadata["rework_attributions"]
    (i.e. its re-execution cost is attributable as review-retry cost).

Hermetic: REDIS_FAKE, no model key, no callback url.

Run from agent_swarm_v6 (install deps first):
    pip install -r orchestrator/requirements.txt
    REDIS_FAKE=1 python scripts/test-usage-cost-phase.py
"""
import asyncio
import json
import os
import sys
from pathlib import Path

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

failures = []


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


async def budget_events(swarm_runtime, swarm_id):
    raw = await swarm_runtime.list_events(swarm_id, limit=500)
    return [e for e in raw["events"] if e["event_type"] == "budget.alert"]


async def main():
    from orchestrator.redis_client import redis_client
    from orchestrator.swarm_runtime import swarm_runtime
    from orchestrator.task_queue import task_queue
    from orchestrator import main as orch

    await redis_client.connect()
    body = {"mode": "swarm", "orchestration_plan": {"objective": "cost phase"},
            "callback": {"url": "", "subscribed_events": []},
            "metadata": {"manager_deployment_id": "m-cost"}}
    run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="c-cost")
    t1 = await task_queue.create_task(task_id=f"{run.swarm_id}-t1", description="impl",
                                      agent_role="impl", required_capabilities=["python"], enqueue=True)
    await swarm_runtime.attach_task(run, t1.task_id)

    usage_result = {"usage": {"model_id": "gpt-x", "model_tokens": 100, "model_cost_usd": 0.01,
                              "billing_source": "newapi"}}

    # 1) Initial execution: no rework attribution yet → cost_phase=initial.
    await orch.emit_usage_event(run, t1, "agent-1", usage_result)
    evs = await budget_events(swarm_runtime, run.swarm_id)
    check("initial usage tagged cost_phase=initial", evs and evs[-1]["payload"].get("cost_phase") == "initial")
    check("usage payload carries attempt", "attempt" in evs[-1]["payload"])

    # 2) Cross-review reopened t1 for rework → recorded in rework_attributions.
    run.metadata["rework_attributions"] = [{"target_task_id": t1.task_id, "root_cause": "test"}]
    await swarm_runtime.save_run(run)
    t1.retry_count = 1
    await orch.emit_usage_event(run, t1, "agent-1", usage_result)
    evs = await budget_events(swarm_runtime, run.swarm_id)
    check("redo usage tagged cost_phase=review_retry", evs[-1]["payload"].get("cost_phase") == "review_retry")
    check("review_retry usage attributable by swarm_id + cost_phase",
          any(e["payload"].get("cost_phase") == "review_retry" and e["payload"].get("swarm_id") == run.swarm_id
              for e in evs))

    # A non-reworked task stays initial even alongside the reworked one.
    t2 = await task_queue.create_task(task_id=f"{run.swarm_id}-t2", description="doc",
                                      agent_role="documentation", required_capabilities=["technical-writing"], enqueue=True)
    await swarm_runtime.attach_task(run, t2.task_id)
    await orch.emit_usage_event(run, t2, "agent-2", usage_result)
    evs = await budget_events(swarm_runtime, run.swarm_id)
    t2_ev = [e for e in evs if e["payload"].get("task_id") == t2.task_id][-1]
    check("non-reworked task stays cost_phase=initial", t2_ev["payload"].get("cost_phase") == "initial")

    # --- /metrics run-level cost_by_phase rollup (#37) ---
    # Persist task results with usage so build_runtime_metrics can split them; t1 is a rework target.
    await task_queue.complete_task(t1.task_id, result=json.dumps({"usage": {"model_cost_usd": 0.05, "model_tokens": 500}}))
    await task_queue.complete_task(t2.task_id, result=json.dumps({"usage": {"model_cost_usd": 0.02, "model_tokens": 200}}))
    refreshed = await swarm_runtime.get_run(run.swarm_id)
    metrics = await orch.build_runtime_metrics(refreshed, "15m", "60s")
    cbp = metrics.get("cost_by_phase") or {}
    check("/metrics exposes cost_by_phase", set(cbp.keys()) == {"initial", "review_retry"})
    check("/metrics review_retry cost = reworked task (t1=0.05)", abs(cbp["review_retry"]["cost_usd"] - 0.05) < 1e-9)
    check("/metrics initial cost = non-reworked task (t2=0.02)", abs(cbp["initial"]["cost_usd"] - 0.02) < 1e-9)
    check("/metrics review_retry tokens = 500", cbp["review_retry"]["model_tokens"] == 500)

    print()
    if failures:
        print(f"{len(failures)} usage cost-phase check(s) FAILED: {failures}")
        sys.exit(1)
    print("all usage cost-phase checks passed")


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