"""M16 bounded soak.

A realistic mixed workload held for a bounded period, measuring what actually drifts: latency
distribution, cold versus warm execution, memory, GPU, database growth, queue depth, leases and
orphaned work. Every number reported here is measured; nothing is extrapolated from a short run.

The workload deliberately mixes bursts, steady state and idle periods. A soak that only hammers
never exercises the keep-warm and unload paths, and a soak that only idles never exercises
admission control.

Production is read, never mutated: no promotion, no alias change, no external workload is touched.
The only writes are gateway requests through a disposable LAB client and the operational polls the
platform already performs for itself.

    python scripts/m16_soak.py --minutes 45 --report soak.json
"""

from __future__ import annotations

import argparse
import json
import os
import random
import socket
import statistics
import subprocess  # noqa: S404 - fixed argv container inspection, never a shell
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from collections import Counter
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from sqlalchemy import create_engine, text
from sqlalchemy.orm import Session

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "backend" / "src"))

from modelforge_api.services.invariants import (  # noqa: E402
    InvariantStatus,
    check_invariants,
)

DEFAULT_BASE_URL = "http://127.0.0.1:8000"
DEFAULT_DATABASE_URL = os.getenv("MODELFORGE_RUNTIME_DATABASE_URL")

SENTENCES = [
    "ModelForge binds applications to capabilities rather than to model paths.",
    "A verified backup is the only backup that counts as a recovery point.",
    "Stale telemetry must never be presented as present-day truth.",
    "Every production promotion carries evidence from local validation.",
    "Downloaded model repositories stay untrusted until supply-chain checks complete.",
    "The scheduler refuses unsafe work instead of risking an out-of-memory failure.",
    "Recovery reconciliation clears current truth and re-measures it.",
    "An embedding-space change never silently reuses incompatible vectors.",
]


@dataclass
class Sample:
    at: float
    kind: str
    ok: bool
    status: int
    latency_ms: float
    detail: dict[str, Any] = field(default_factory=dict)


@dataclass
class SoakState:
    samples: list[Sample] = field(default_factory=list)
    lock: threading.Lock = field(default_factory=threading.Lock)
    stop: threading.Event = field(default_factory=threading.Event)

    def record(self, sample: Sample) -> None:
        with self.lock:
            self.samples.append(sample)


def docker(*args: str) -> str:
    try:
        completed = subprocess.run(  # noqa: S603 - fixed argv, no shell
            ["docker", *args],  # noqa: S607
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=60,
            check=False,
        )
    except (OSError, subprocess.TimeoutExpired):
        return ""
    return (completed.stdout or "").strip()


def container_stats(names: list[str]) -> dict[str, dict[str, Any]]:
    """Point-in-time RSS and CPU per container, read from the runtime rather than guessed."""

    raw = docker(
        "stats",
        "--no-stream",
        "--format",
        "{{.Name}}\t{{.MemUsage}}\t{{.CPUPerc}}\t{{.PIDs}}",
        *names,
    )
    result: dict[str, dict[str, Any]] = {}
    for line in raw.splitlines():
        parts = line.split("\t")
        if len(parts) != 4:
            continue
        name, memory, cpu, pids = parts
        used = memory.split("/")[0].strip()
        result[name] = {"memory": used, "cpu_percent": cpu.strip(), "pids": pids.strip()}
    return result


def restart_counts(names: list[str]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for name in names:
        value = docker("inspect", "--format", "{{.RestartCount}}", name)
        counts[name] = int(value) if value.isdigit() else -1
    return counts


def http(
    url: str,
    *,
    token: str | None = None,
    bearer: str | None = None,
    method: str = "GET",
    payload: dict[str, Any] | None = None,
    timeout: float = 180.0,
) -> tuple[bool, int, float, dict[str, Any]]:
    body = json.dumps(payload).encode() if payload is not None else None
    request = urllib.request.Request(url, data=body, method=method)  # noqa: S310 - fixed scheme
    if token:
        request.add_header("X-ModelForge-Admin-Token", token)
    if bearer:
        request.add_header("Authorization", f"Bearer {bearer}")
    if body is not None:
        request.add_header("Content-Type", "application/json")
    started = time.perf_counter()
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:  # noqa: S310
            raw = response.read()
            elapsed = (time.perf_counter() - started) * 1000
            try:
                decoded = json.loads(raw.decode("utf-8"))
            except (UnicodeDecodeError, json.JSONDecodeError):
                decoded = {}
            return True, response.status, elapsed, decoded if isinstance(decoded, dict) else {}
    except urllib.error.HTTPError as error:
        elapsed = (time.perf_counter() - started) * 1000
        try:
            decoded = json.loads(error.read().decode("utf-8"))
        except Exception:  # noqa: BLE001 - an unreadable error body is still an error
            decoded = {}
        return False, error.code, elapsed, decoded if isinstance(decoded, dict) else {}
    except (urllib.error.URLError, TimeoutError, OSError) as error:
        elapsed = (time.perf_counter() - started) * 1000
        return False, 0, elapsed, {"transport_error": type(error).__name__}


# The soak reports latency percentiles, so the client-side floor has to be negligible before any
# of them mean anything. A first run measured a p50 of two seconds on every route including
# /health/live; the cause was the harness resolving "localhost" to ::1 on Windows and waiting out
# the connect timeout before falling back to IPv4. That inflated every percentile by a constant
# ~2,045 ms of client cost and said nothing about the platform. Measure the floor and refuse rather
# than publish a number that is mostly the measuring instrument.
MAX_CONNECT_FLOOR_MS = 100.0


def connect_floor_ms(base_url: str, samples: int = 7) -> float:
    """Median TCP connect time to the target, excluding anything the platform does."""

    parsed = urllib.parse.urlparse(base_url)
    host = parsed.hostname or "127.0.0.1"
    port = parsed.port or (443 if parsed.scheme == "https" else 80)
    timings = []
    for _ in range(samples):
        started = time.perf_counter()
        try:
            connection = socket.create_connection((host, port), timeout=10)
        except OSError:
            return float("inf")
        connection.close()
        timings.append((time.perf_counter() - started) * 1000)
    return statistics.median(timings)


# --------------------------------------------------------------------- workers


def gateway_worker(
    state: SoakState, base_url: str, secret: str, rng: random.Random, worker_id: int
) -> None:
    """Real capability traffic with bursts, steady state and idle gaps."""

    while not state.stop.is_set():
        phase = rng.choices(["burst", "steady", "idle"], weights=[2, 5, 3])[0]
        if phase == "idle":
            state.stop.wait(rng.uniform(8, 25))
            continue
        count = rng.randint(3, 6) if phase == "burst" else 1
        for _ in range(count):
            if state.stop.is_set():
                return
            batch = [rng.choice(SENTENCES) for _ in range(rng.randint(1, 4))]
            ok, status, latency, body = http(
                f"{base_url}/api/v1/capabilities/rag.embedding@1/invoke",
                bearer=secret,
                method="POST",
                payload={"input": batch},
            )
            execution = body.get("execution") or {}
            timings = execution.get("timings") or {}
            state.record(
                Sample(
                    at=time.time(),
                    kind="gateway.rag.embedding",
                    ok=ok,
                    status=status,
                    latency_ms=latency,
                    detail={
                        "cold": bool(execution.get("cold")),
                        "residency": execution.get("residency"),
                        "queue_ms": timings.get("queue_ms"),
                        "inference_ms": timings.get("inference_ms"),
                        "load_ms": timings.get("load_ms"),
                        "worker": worker_id,
                        "batch": len(batch),
                        "code": (body.get("error") or {}).get("code") if not ok else None,
                    },
                )
            )
            if phase == "burst":
                state.stop.wait(rng.uniform(0.05, 0.4))
            else:
                state.stop.wait(rng.uniform(1.5, 5.0))


READ_ROUTES = [
    ("registry.models", "/api/v1/models", False),
    ("registry.projects", "/api/v1/projects", False),
    ("hardware", "/api/v1/hardware", False),
    ("health.ready", "/api/v1/health/ready", False),
    ("operations.overview", "/api/v1/admin/operations/overview", True),
    ("operations.alerts", "/api/v1/admin/operations/alerts?limit=50", True),
    ("lifecycle.operations", "/api/v1/admin/lifecycle/operations", True),
    ("migration.plans", "/api/v1/admin/migrations/plans", True),
    ("recovery.dashboard", "/api/v1/admin/recovery/dashboard", True),
    ("recovery.backups", "/api/v1/admin/recovery/backups?limit=20", True),
]


def read_worker(state: SoakState, base_url: str, token: str, rng: random.Random) -> None:
    while not state.stop.is_set():
        kind, path, needs_token = rng.choice(READ_ROUTES)
        ok, status, latency, _ = http(
            f"{base_url}{path}", token=token if needs_token else None, timeout=60.0
        )
        state.record(Sample(at=time.time(), kind=kind, ok=ok, status=status, latency_ms=latency))
        state.stop.wait(rng.uniform(0.5, 3.0))


def background_worker(state: SoakState, base_url: str, token: str) -> None:
    """The operational work the platform performs for itself, driven at a bounded rate."""

    while not state.stop.is_set():
        for kind, path in (
            ("ops.capacity_collect", "/api/v1/admin/operations/capacity/collect"),
            ("ops.slo_evaluate", "/api/v1/admin/operations/slo-evaluations/run"),
            ("ops.alert_evaluate", "/api/v1/admin/operations/alerts/evaluate"),
        ):
            if state.stop.is_set():
                return
            ok, status, latency, _ = http(
                f"{base_url}{path}", token=token, method="POST", timeout=120.0
            )
            state.record(
                Sample(at=time.time(), kind=kind, ok=ok, status=status, latency_ms=latency)
            )
        state.stop.wait(60)


# --------------------------------------------------------------------- measurement


def resource_snapshot(database_url: str, containers: list[str]) -> dict[str, Any]:
    engine = create_engine(database_url, pool_pre_ping=True)
    try:
        with Session(engine) as session:
            counts = {
                table: int(
                    session.execute(text(f"select count(*) from {table}")).scalar_one()  # noqa: S608
                )
                for table in (
                    "audit_events",
                    "gateway_requests",
                    "serving_jobs",
                    "serving_gpu_leases",
                    "slo_evaluations",
                    "capacity_snapshots",
                    "alert_history_events",
                    "operational_alerts",
                    "placement_plans",
                    "hardware_inventory_runs",
                )
            }
            active_leases = int(
                session.execute(
                    text(
                        "select count(*) from serving_gpu_leases "
                        "where released_at is null and state in "
                        "('pending','granted','active','held')"
                    )
                ).scalar_one()
            )
            in_flight = int(
                session.execute(
                    text(
                        "select count(*) from serving_jobs "
                        "where status in ('queued','leased','running')"
                    )
                ).scalar_one()
            )
            residency = int(
                session.execute(
                    text(
                        "select count(*) from residency_allocations "
                        "where state in ('loading','resident','unloading')"
                    )
                ).scalar_one()
            )
            connections = int(
                session.execute(
                    text(
                        "select count(*) from pg_stat_activity "
                        "where datname = current_database()"
                    )
                ).scalar_one()
            )
            gpu = session.execute(
                text(
                    "select a.name, t.used_vram_bytes, t.free_vram_bytes, "
                    "t.gpu_utilization_percent, t.observed_at "
                    "from accelerator_telemetry_latest t "
                    "join accelerators a on a.id = t.accelerator_id"
                )
            ).all()
            scheduler = session.execute(
                text(
                    "select pressure_state, recovery_candidate, last_observed_at "
                    "from scheduler_accelerator_states"
                )
            ).all()
            capacity = session.execute(
                text(
                    "select gpu_external_bytes, gpu_schedulable_bytes, gpu_managed_resident_bytes, "
                    "pressure_state, observed_at from capacity_snapshots "
                    "order by observed_at desc limit 1"
                )
            ).first()
    finally:
        engine.dispose()
    return {
        "observed_at": datetime.now(UTC).isoformat(),
        "row_counts": counts,
        "active_leases": active_leases,
        "serving_jobs_in_flight": in_flight,
        "residency_allocations": residency,
        "database_connections": connections,
        "gpu": [
            {
                "name": name,
                "used_vram_bytes": used,
                "free_vram_bytes": free,
                "utilisation_percent": util,
                "observed_at": observed.isoformat() if observed else None,
            }
            for name, used, free, util, observed in gpu
        ],
        "scheduler": [
            {
                "pressure_state": pressure,
                "recovery_candidate": bool(candidate),
                "last_observed_at": observed.isoformat() if observed else None,
            }
            for pressure, candidate, observed in scheduler
        ],
        "latest_capacity": (
            {
                "gpu_external_bytes": capacity[0],
                "gpu_schedulable_bytes": capacity[1],
                "gpu_managed_resident_bytes": capacity[2],
                "pressure_state": capacity[3],
                "observed_at": capacity[4].isoformat() if capacity[4] else None,
            }
            if capacity
            else None
        ),
        "containers": container_stats(containers),
        "restart_counts": restart_counts(containers),
    }


def percentile(values: list[float], fraction: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    index = min(len(ordered) - 1, max(0, round(fraction * (len(ordered) - 1))))
    return round(ordered[index], 3)


def summarise(samples: list[Sample]) -> dict[str, Any]:
    by_kind: dict[str, list[Sample]] = {}
    for sample in samples:
        by_kind.setdefault(sample.kind, []).append(sample)

    result: dict[str, Any] = {}
    for kind, items in sorted(by_kind.items()):
        latencies = [item.latency_ms for item in items if item.ok]
        statuses = Counter(item.status for item in items)
        failures = [item for item in items if not item.ok]
        capacity = [
            item
            for item in failures
            if item.status in (429, 503) or item.detail.get("code") in ("QUEUE_FULL", "RATE_LIMITED")
        ]
        internal = [item for item in failures if item.status >= 500 or item.status == 0]
        entry: dict[str, Any] = {
            "requests": len(items),
            "succeeded": len(items) - len(failures),
            "capacity_rejected": len(capacity),
            "internal_errors": len(internal),
            "statuses": dict(sorted(statuses.items())),
            "p50_ms": percentile(latencies, 0.50),
            "p95_ms": percentile(latencies, 0.95),
            "p99_ms": percentile(latencies, 0.99),
            "mean_ms": round(statistics.fmean(latencies), 3) if latencies else None,
            "max_ms": round(max(latencies), 3) if latencies else None,
        }
        if kind.startswith("gateway."):
            cold = [item for item in items if item.ok and item.detail.get("cold")]
            warm = [item for item in items if item.ok and not item.detail.get("cold")]
            queue = [
                float(item.detail["queue_ms"])
                for item in items
                if item.ok and item.detail.get("queue_ms") is not None
            ]
            inference = [
                float(item.detail["inference_ms"])
                for item in items
                if item.ok and item.detail.get("inference_ms") is not None
            ]
            entry["cold_starts"] = len(cold)
            entry["warm_invocations"] = len(warm)
            entry["queue_p95_ms"] = percentile(queue, 0.95)
            entry["inference_p50_ms"] = percentile(inference, 0.50)
            entry["inference_p95_ms"] = percentile(inference, 0.95)
            entry["failure_codes"] = dict(
                Counter(
                    item.detail.get("code") for item in failures if item.detail.get("code")
                )
            )
        result[kind] = entry
    return result


# --------------------------------------------------------------------- runner


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
    parser.add_argument("--database-url", default=DEFAULT_DATABASE_URL)
    parser.add_argument("--token", default=None)
    parser.add_argument("--secret", default=None, help="gateway client secret")
    parser.add_argument("--minutes", type=float, default=45.0)
    parser.add_argument("--gateway-workers", type=int, default=3)
    parser.add_argument("--read-workers", type=int, default=2)
    parser.add_argument("--seed", type=int, default=20260827)
    parser.add_argument("--report", default=None)
    parser.add_argument("--progress-seconds", type=float, default=120.0)
    args = parser.parse_args(argv)
    if not args.database_url:
        parser.error(
            "--database-url or MODELFORGE_RUNTIME_DATABASE_URL is required "
            "(use the non-owner runtime role)"
        )

    root = Path(__file__).resolve().parents[1]
    token = args.token
    if not token and (root / ".env").is_file():
        for line in (root / ".env").read_text("utf-8").splitlines():
            if line.startswith("MODELFORGE_OPERATOR_API_KEY="):
                token = line.split("=", 1)[1].strip()
    if not token:
        parser.error("an operator API key is required")
    if not args.secret:
        parser.error("a gateway client secret is required")

    containers = [
        "modelforge-api-1",
        "modelforge-postgres-1",
        "modelforge-redis-1",
        "modelforge-web-1",
    ]

    floor = connect_floor_ms(args.base_url)
    print(f"client connect floor: {floor:.1f} ms", flush=True)
    if floor > MAX_CONNECT_FLOOR_MS:
        print(
            f"refusing to soak: the client connect floor is {floor:.1f} ms, above the "
            f"{MAX_CONNECT_FLOOR_MS:.0f} ms budget. Every latency percentile would be dominated "
            "by the harness rather than the platform. Use an address that resolves directly "
            "(127.0.0.1 rather than localhost).",
            flush=True,
        )
        return 2

    print(f"soak: {args.minutes:.1f} minutes, seed {args.seed}", flush=True)
    before = resource_snapshot(args.database_url, containers)
    engine = create_engine(args.database_url, pool_pre_ping=True)
    with Session(engine) as session:
        baseline = check_invariants(session)
    engine.dispose()
    if baseline.violated:
        print("refusing to soak: invariants are already violated", flush=True)
        return 2

    state = SoakState()
    threads: list[threading.Thread] = []
    for index in range(args.gateway_workers):
        thread = threading.Thread(
            target=gateway_worker,
            args=(state, args.base_url, args.secret, random.Random(args.seed + index), index),
            daemon=True,
        )
        threads.append(thread)
    for index in range(args.read_workers):
        thread = threading.Thread(
            target=read_worker,
            args=(state, args.base_url, token, random.Random(args.seed + 100 + index)),
            daemon=True,
        )
        threads.append(thread)
    threads.append(
        threading.Thread(
            target=background_worker, args=(state, args.base_url, token), daemon=True
        )
    )

    started = time.time()
    deadline = started + args.minutes * 60
    for thread in threads:
        thread.start()

    invariant_checks: list[dict[str, Any]] = []
    progress: list[dict[str, Any]] = []
    try:
        next_progress = started + args.progress_seconds
        while time.time() < deadline:
            time.sleep(2)
            if time.time() >= next_progress:
                next_progress = time.time() + args.progress_seconds
                engine = create_engine(args.database_url, pool_pre_ping=True)
                with Session(engine) as session:
                    report = check_invariants(session)
                engine.dispose()
                with state.lock:
                    total = len(state.samples)
                    gateway = sum(
                        1 for item in state.samples if item.kind.startswith("gateway.")
                    )
                    failed = sum(1 for item in state.samples if not item.ok)
                invariant_checks.append(
                    {
                        "at": datetime.now(UTC).isoformat(),
                        "violated": report.violated,
                        "violations": [
                            item.key
                            for item in report.results
                            if item.status is InvariantStatus.VIOLATED
                        ],
                    }
                )
                elapsed = (time.time() - started) / 60
                entry = {
                    "minutes": round(elapsed, 1),
                    "requests": total,
                    "gateway": gateway,
                    "failed": failed,
                    "invariants_violated": report.violated,
                }
                progress.append(entry)
                print(f"  {json.dumps(entry)}", flush=True)
                if report.violated:
                    print("  invariant violated; stopping the soak early", flush=True)
                    break
    except KeyboardInterrupt:
        print("  interrupted; reporting partial evidence", flush=True)
    finally:
        state.stop.set()
        for thread in threads:
            thread.join(timeout=240)

    duration = time.time() - started
    after = resource_snapshot(args.database_url, containers)
    engine = create_engine(args.database_url, pool_pre_ping=True)
    with Session(engine) as session:
        final = check_invariants(session)
    engine.dispose()

    with state.lock:
        samples = list(state.samples)

    drift = {
        table: after["row_counts"][table] - before["row_counts"][table]
        for table in before["row_counts"]
    }
    report = {
        "generated_at": datetime.now(UTC).isoformat(),
        "seed": args.seed,
        "client_connect_floor_ms": round(floor, 3),
        "requested_minutes": args.minutes,
        "actual_duration_seconds": round(duration, 1),
        "gateway_workers": args.gateway_workers,
        "read_workers": args.read_workers,
        "total_requests": len(samples),
        "total_failures": sum(1 for item in samples if not item.ok),
        "by_kind": summarise(samples),
        "resources_before": before,
        "resources_after": after,
        "row_growth": drift,
        "invariant_checks": invariant_checks,
        "invariants_final": {
            "checked": final.checked,
            "violated": final.violated,
            "violations": [
                {"key": item.key, "examples": item.violations}
                for item in final.results
                if item.status is InvariantStatus.VIOLATED
            ],
        },
        "progress": progress,
    }

    print(json.dumps({key: report[key] for key in ("actual_duration_seconds", "total_requests", "total_failures")}, indent=2))
    print(json.dumps(report["by_kind"], indent=2))
    print("row growth:", json.dumps(drift))
    print("leases after:", after["active_leases"], "| jobs in flight:", after["serving_jobs_in_flight"])
    print("invariants:", final.holding, "/", final.checked, "hold")

    if args.report:
        Path(args.report).write_text(json.dumps(report, indent=2), encoding="utf-8")
        print(f"report written to {args.report}")
    return 0 if final.violated == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main())