"""Population simulator. All state changes go through the public HTTP API."""
from __future__ import annotations

import argparse
import random
from dataclasses import dataclass

import httpx


@dataclass
class Experiment:
    clients: int = 100
    runs_per_client: int = 10
    mode: str = "diversified"
    strategy: str = "balanced"
    bad: tuple[str, ...] = ()
    bad_pair: tuple[str, str] | None = None
    seed: int = 7


def run_experiment(base_url: str, release_id: int, cfg: Experiment, transport=None):
    rng=random.Random(cfg.seed); registered=[]; failures=0
    with httpx.Client(base_url=base_url,transport=transport,timeout=30) as api:
        first=None
        for n in range(cfg.clients):
            identity=f"sim-{cfg.mode}-{cfg.strategy}-{cfg.seed}-{n}"
            strategy="homogeneous" if cfg.mode == "homogeneous" else cfg.strategy
            comp=api.post("/api/v1/clients",json={"release_id":release_id,"installation_id":identity,"strategy":strategy}).raise_for_status().json()
            if cfg.mode == "homogeneous": first=comp
            registered.append(comp)
        for comp in registered:
            active=set(comp["implementations"].values())
            direct_faults=active.intersection(cfg.bad)
            pair_faults=set(cfg.bad_pair) if cfg.bad_pair and set(cfg.bad_pair)<=active else set()
            fault_ids=direct_faults | pair_faults
            faulty=bool(fault_ids)
            for _ in range(cfg.runs_per_client):
                failed=faulty and rng.random()<0.75
                calls=[{"spec_function":k,"implementation_id":v,
                        "outcome":"error" if failed and v in fault_ids else "ok"}
                       for k,v in comp["implementations"].items()]
                ex=api.post("/api/v1/executions",json={"client_id":comp["client_id"],"operation":"extract","success":not failed,
                    "composition_revision":comp["revision"],
                    "error_kind":"controlled_fault" if failed else None,
                    "calls":calls}).raise_for_status().json()
                if failed:
                    failures+=1
                    focal_candidates=cfg.bad_pair if pair_faults else cfg.bad
                    impl=next(x for x in focal_candidates if x in fault_ids)
                    spec=next(k for k,v in comp["implementations"].items() if v==impl)
                    api.post("/api/v1/telemetry",json={"client_id":comp["client_id"],"execution_id":ex["execution_id"],
                        "release_id":comp["release_id"],"composition_revision":comp["revision"],"event_type":"extraction_failure",
                        "spec_function":spec,"implementation_id":impl,"message_code":"controlled_simulator_failure",
                        "recent_calls":calls[-32:]}).raise_for_status()
        return {"mode":cfg.mode,"clients":cfg.clients,"executions":cfg.clients*cfg.runs_per_client,
                "failures":failures,"affected_fraction":sum(bool(set(c["implementations"].values()).intersection(cfg.bad)) or bool(cfg.bad_pair and set(cfg.bad_pair)<=set(c["implementations"].values())) for c in registered)/len(registered),
                "summary":api.get("/api/v1/dashboard/summary",params={"release_id":release_id}).raise_for_status().json()}


def run_quarantine_experiment(base_url: str, release_id: int, cfg: Experiment,
                              quarantine_ids: tuple[str, ...], transport=None):
    """Run the same client cohort before and after server-side quarantine."""
    before=run_experiment(base_url,release_id,cfg,transport)
    changes=[]
    with httpx.Client(base_url=base_url,transport=transport,timeout=30) as api:
        for impl_id in quarantine_ids:
            change=api.put(f"/api/v1/implementations/{impl_id}/status",params={"release_id":release_id},json={
                "status":"quarantined","reason":"population experiment"
            }).raise_for_status().json()
            changes.append(change)
    after=run_experiment(base_url,release_id,cfg,transport)
    before_rate=before["failures"]/before["executions"] if before["executions"] else 0.0
    after_rate=(after["failures"]/after["executions"] if after["executions"] else 0.0)
    return {"before":before,"quarantine":changes,"after":after,
            "failure_rate_drop":before_rate-after_rate}


def main():
    p=argparse.ArgumentParser(); p.add_argument("--url",default="http://127.0.0.1:8000"); p.add_argument("--release",type=int,required=True)
    p.add_argument("--clients",type=int,default=100); p.add_argument("--runs",type=int,default=10)
    p.add_argument("--mode",choices=["homogeneous","diversified"],default="diversified"); p.add_argument("--strategy",choices=["balanced","random"],default="balanced")
    p.add_argument("--bad",action="append",default=[]); p.add_argument("--bad-pair",nargs=2)
    a=p.parse_args(); print(run_experiment(a.url,a.release,Experiment(a.clients,a.runs,a.mode,a.strategy,tuple(a.bad),tuple(a.bad_pair) if a.bad_pair else None)))

if __name__ == "__main__": main()
