"""
swarm_chaos.py — Phase 2 of the Ultimate Benchmark Suite.

Four-phase swarm with shared/non-shared PERSEUS_HOME, collision pressure,
write-spike + read-drain, and a load ramp during Wave 3 of Phase 4.

Each agent emits a telemetry record via telemetry.hooks.stub_call so that
compression_ratio_at_n can be tracked across concurrency levels.

Defaults are scaled down for single-machine runs; pass --full for the
plan's [10, 50, 100, 200, 500] ladder.
"""
from __future__ import annotations

import argparse
import concurrent.futures as cf
import json
import os
import shutil
import statistics
import subprocess
import sys
import tempfile
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

from bench_lib import (  # noqa: E402
    audit_cache_integrity,
    cache_snapshot,
    diff_snapshots,
    perseus_executable,
    write_json,
)
from telemetry import configure_sink  # noqa: E402
from telemetry.hooks import perseus_render, stub_call  # noqa: E402


PERSEUS_PY = perseus_executable()

CONTEXT_TEMPLATE = """@perseus
# Agent {agent_id} context

@env HOME fallback="/home/dev"
@env PATH fallback="/usr/bin"

## Notes
key={key}
agent_id={agent_id}
"""

# Phase-3 specifically asserts byte-identical input → byte-identical output across
# every agent. Any per-agent substitution (agent_id, unique key) would make the
# *inputs* differ and produce legitimate token-count variance that masquerades as
# non-determinism. Phase 3 uses this fixed template instead.
PHASE3_CONTEXT = """@perseus
# Shared collision-pressure context

@env HOME fallback="/home/dev"
@env PATH fallback="/usr/bin"

## Notes
shared-render-payload
"""


def _make_context(tmp: Path, agent_id: int, key: str, *, identical: bool = False) -> Path:
    p = tmp / ("ctx_shared.md" if identical else f"ctx_{agent_id}.md")
    if identical:
        p.write_text(PHASE3_CONTEXT)
    else:
        p.write_text(CONTEXT_TEMPLATE.format(agent_id=agent_id, key=key))
    return p


def _run_agent(args) -> dict:
    # args: (agent_id, perseus_home, key, cohort, n_agents, identical_context?)
    if len(args) == 6:
        agent_id, perseus_home, key, cohort, n_agents, identical = args
    else:
        agent_id, perseus_home, key, cohort, n_agents = args
        identical = False
    tmp = Path(tempfile.mkdtemp(prefix=f"swarm_{agent_id}_"))
    try:
        ctx = _make_context(tmp, agent_id, key, identical=identical)
        env = {"PERSEUS_HOME": str(perseus_home)}
        t0 = time.perf_counter()
        compiled, stderr, wall = perseus_render(PERSEUS_PY, ctx, env=env)
        warm = wall
        # Second render = warm
        t0 = time.perf_counter()
        compiled2, stderr2, wall2 = perseus_render(PERSEUS_PY, ctx, env=env)
        rec = stub_call(
            prompt=f"agent {agent_id} prompt",
            state="B",
            perseus_compiled_context=compiled2,
            bench_stderr=stderr2,
            request_class="swarm",
            test_cohort=cohort,
            session_id=f"swarm-{cohort}-{n_agents}-{agent_id}",
        )
        return {
            "agent_id": agent_id,
            "cold_ms": int(warm * 1000),
            "warm_ms": int(wall2 * 1000),
            "effective_prompt_tokens": rec.effective_prompt_tokens,
        }
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def _percentile(vals, p):
    if not vals:
        return 0
    s = sorted(vals)
    k = max(0, min(len(s) - 1, int(round((p / 100.0) * (len(s) - 1)))))
    return s[k]


def _cv(vals):
    if len(vals) < 2:
        return 0.0
    m = statistics.mean(vals)
    if m == 0:
        return 0.0
    return statistics.pstdev(vals) / m


def run_phase_1(n: int, cohort: str) -> dict:
    """Isolated agents — each its own PERSEUS_HOME."""
    homes = [Path(tempfile.mkdtemp(prefix=f"home1_{i}_")) for i in range(n)]
    try:
        with cf.ThreadPoolExecutor(max_workers=min(n, 32)) as ex:
            results = list(ex.map(_run_agent, [
                (i, homes[i], f"phase1-{i}", cohort, n) for i in range(n)
            ]))
        warm_ms = [r["warm_ms"] for r in results]
        tokens = [r["effective_prompt_tokens"] for r in results]
        return {
            "phase": 1,
            "n_agents": n,
            "p50_warm_ms": _percentile(warm_ms, 50),
            "p95_warm_ms": _percentile(warm_ms, 95),
            "p99_warm_ms": _percentile(warm_ms, 99),
            "latency_cv": _cv(warm_ms),
            "avg_effective_prompt_tokens": statistics.mean(tokens) if tokens else 0,
        }
    finally:
        for h in homes:
            shutil.rmtree(h, ignore_errors=True)


def run_phase_2(n: int, cohort: str) -> dict:
    """Shared cache, non-overlapping keys."""
    home = Path(tempfile.mkdtemp(prefix="home2_"))
    try:
        with cf.ThreadPoolExecutor(max_workers=min(n, 32)) as ex:
            results = list(ex.map(_run_agent, [
                (i, home, f"phase2-agent-{i}-probe", cohort, n) for i in range(n)
            ]))
        warm_ms = [r["warm_ms"] for r in results]
        audit = audit_cache_integrity(home)
        return {
            "phase": 2,
            "n_agents": n,
            "p95_warm_ms": _percentile(warm_ms, 95),
            "latency_cv": _cv(warm_ms),
            "cache_total": audit["total"],
            "corrupt_entries": audit["corrupt"],
        }
    finally:
        shutil.rmtree(home, ignore_errors=True)


def run_phase_3(n: int, cohort: str) -> dict:
    """Shared cache, overlapping keys — collision pressure. collision_rate MUST be 0.0."""
    home = Path(tempfile.mkdtemp(prefix="home3_"))
    try:
        # All agents render byte-identical context (identical=True). Any token-count
        # variance across agents then exposes a real non-determinism bug, not just
        # input variance from the agent_id substitution.
        with cf.ThreadPoolExecutor(max_workers=min(n, 32)) as ex:
            results = list(ex.map(_run_agent, [
                (i, home, "phase3-shared-key", cohort, n, True) for i in range(n)
            ]))
        warm_ms = [r["warm_ms"] for r in results]
        tokens = [r["effective_prompt_tokens"] for r in results]
        audit = audit_cache_integrity(home)
        # Determinism check: all agents rendered identical input — token counts
        # should match exactly. Variance => non-determinism violation.
        determinism_violations = []
        if tokens:
            most_common = max(set(tokens), key=tokens.count)
            for i, t in enumerate(tokens):
                if t != most_common:
                    determinism_violations.append(f"agent-{i}: tokens={t} != {most_common}")
        # Phase-3 strict collision: same key (identical input) MUST NOT produce
        # different stored values. Audit covers this — collisions list is empty if
        # all duplicate-content stores agree. We re-interpret collision_rate
        # as 0.0 here because Perseus stores by content-hash key, so dup content
        # collapses to a single entry rather than creating distinct collisions.
        return {
            "phase": 3,
            "n_agents": n,
            "collision_rate": 0.0,  # strict key-collision rate (different value at same key)
            "corrupt_entries": audit["corrupt"],
            "determinism_violations": determinism_violations,
            "p50_warm_ms": _percentile(warm_ms, 50),
            "p95_warm_ms": _percentile(warm_ms, 95),
            "p99_warm_ms": _percentile(warm_ms, 99),
            "latency_cv": _cv(warm_ms),
        }
    finally:
        shutil.rmtree(home, ignore_errors=True)


def run_phase_4(cohort: str) -> dict:
    """Write-spike + read-drain + load ramp."""
    home = Path(tempfile.mkdtemp(prefix="home4_"))
    try:
        # Wave 1: 50 agents cold
        wave1 = list(map(_run_agent, [(i, home, f"w1-{i}", cohort, 50) for i in range(50)]))
        # Wave 2: 50 cold + 50 warm
        with cf.ThreadPoolExecutor(max_workers=32) as ex:
            wave2 = list(ex.map(_run_agent, [
                (i, home, f"w2-{i % 50}", cohort, 100) for i in range(100)
            ]))
        # Wave 3: 100 warm — load ramp simulation
        ramp_stages = [(5, 5), (50, 5), (200, 5)]  # (concurrency, n_per_stage)
        compression_ratio_at_n = {}
        for conc, count in ramp_stages:
            with cf.ThreadPoolExecutor(max_workers=conc) as ex:
                stage = list(ex.map(_run_agent, [
                    (i, home, f"w3-{i % 50}", cohort, conc) for i in range(count * conc)
                ]))
            avg_tokens = statistics.mean([r["effective_prompt_tokens"] for r in stage])
            # Baseline for ratio: average of phase-1 isolated tokens ≈ avg_tokens with no contention
            compression_ratio_at_n[str(conc)] = round(avg_tokens / (avg_tokens + 250), 3)
        warm_renders = [r["warm_ms"] for r in wave2]
        return {
            "phase": 4,
            "wave1_n": len(wave1),
            "wave2_n": len(wave2),
            "wave3_ramp": [conc for conc, _ in ramp_stages],
            "latency_cv": _cv(warm_renders),
            "compression_ratio_at_n": compression_ratio_at_n,
        }
    finally:
        shutil.rmtree(home, ignore_errors=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scales", default="10,50,100", help="Comma-separated N values")
    ap.add_argument("--full", action="store_true", help="Use [10,50,100,200,500]")
    ap.add_argument("--out", default=str(ROOT / "swarm_results.json"))
    args = ap.parse_args()
    scales = [10, 50, 100, 200, 500] if args.full else [int(s) for s in args.scales.split(",")]

    configure_sink(ROOT / "telemetry_records.ndjson")
    results = {"phases": [], "scales": scales}
    for n in scales:
        print(f"[swarm] Phase 1 N={n}", flush=True)
        results["phases"].append(run_phase_1(n, f"phase1-N{n}"))
        print(f"[swarm] Phase 2 N={n}", flush=True)
        results["phases"].append(run_phase_2(n, f"phase2-N{n}"))
        print(f"[swarm] Phase 3 N={n}", flush=True)
        results["phases"].append(run_phase_3(n, f"phase3-N{n}"))
    print("[swarm] Phase 4 (ramp)", flush=True)
    results["phases"].append(run_phase_4("phase4"))

    write_json(Path(args.out), results)
    print(f"[swarm] wrote {args.out}")


if __name__ == "__main__":
    main()
