#!/usr/bin/env python3
"""Qualify Gazebo pause, single-step, speed and resource behavior on SIM S0."""

from __future__ import annotations

import argparse
import json
import re
import shutil
import subprocess
import threading
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import yaml

TIME_SCHEMA = "missioncore.simulation-time-control/v1"
RESOURCE_SCHEMA = "missioncore.simulation-resource-baseline/v1"
MIB = 1024**2
GIB = 1024**3


class ProbeError(RuntimeError):
    """A runtime observation could not satisfy the reviewed S0 contract."""


@dataclass(frozen=True, slots=True)
class WorldStats:
    sim_time_ns: int
    real_time_ns: int
    iterations: int
    paused: bool
    real_time_factor: float
    step_size_ns: int


def _run(command: list[str], *, timeout: float) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(
        command,
        check=False,
        capture_output=True,
        text=True,
        timeout=timeout,
    )
    if result.returncode != 0:
        raise ProbeError(
            f"command failed ({result.returncode}): {' '.join(command)}; "
            f"stderr={result.stderr.strip()!r}"
        )
    return result


def _time_field(message: str, field: str) -> int:
    match = re.search(rf"\b{re.escape(field)}\s*\{{(?P<body>.*?)\}}", message, re.DOTALL)
    if match is None:
        raise ProbeError(f"Gazebo stats omitted {field}")
    body = match.group("body")
    seconds = re.search(r"\bsec:\s*(-?\d+)", body)
    nanoseconds = re.search(r"\bnsec:\s*(-?\d+)", body)
    sec_value = int(seconds.group(1)) if seconds is not None else 0
    nsec_value = int(nanoseconds.group(1)) if nanoseconds is not None else 0
    return sec_value * 1_000_000_000 + nsec_value


def _scalar_field(
    message: str,
    field: str,
    converter: type[int] | type[float],
    *,
    default: int | float | None = None,
) -> int | float:
    match = re.search(rf"\b{re.escape(field)}:\s*([^\s]+)", message)
    if match is None:
        if default is None:
            raise ProbeError(f"Gazebo stats omitted {field}")
        return default
    return converter(match.group(1))


def _stats_snapshot() -> WorldStats:
    result = _run(
        ["gz", "topic", "--echo", "--topic", "/stats", "-n", "1"],
        timeout=8,
    )
    paused_match = re.search(r"\bpaused:\s*(true|false)", result.stdout)
    paused = paused_match is not None and paused_match.group(1) == "true"
    return WorldStats(
        sim_time_ns=_time_field(result.stdout, "sim_time"),
        real_time_ns=_time_field(result.stdout, "real_time"),
        iterations=int(_scalar_field(result.stdout, "iterations", int)),
        paused=paused,
        real_time_factor=float(
            _scalar_field(result.stdout, "real_time_factor", float, default=0.0)
        ),
        step_size_ns=_time_field(result.stdout, "step_size"),
    )


def _control(world: str, request: str) -> str:
    result = _run(
        [
            "gz",
            "service",
            "-s",
            f"/world/{world}/control",
            "--reqtype",
            "gz.msgs.WorldControl",
            "--reptype",
            "gz.msgs.Boolean",
            "--timeout",
            "5000",
            "--req",
            request,
        ],
        timeout=8,
    )
    if "data: true" not in result.stdout:
        raise ProbeError(
            f"Gazebo rejected world control request {request!r}: {result.stdout.strip()!r}"
        )
    return result.stdout.strip()


def _wait_for_pause_state(expected: bool, *, timeout: float = 6) -> WorldStats:
    deadline = time.monotonic() + timeout
    observed: WorldStats | None = None
    while time.monotonic() < deadline:
        observed = _stats_snapshot()
        if observed.paused is expected:
            return observed
    raise ProbeError(f"Gazebo pause state did not become {expected}; last={observed!r}")


def _wait_for_iteration(minimum: int, *, timeout: float = 6) -> WorldStats:
    deadline = time.monotonic() + timeout
    observed: WorldStats | None = None
    while time.monotonic() < deadline:
        observed = _stats_snapshot()
        if observed.iterations >= minimum:
            return observed
    raise ProbeError(f"Gazebo iteration did not reach {minimum}; last={observed!r}")


def _meminfo() -> dict[str, int]:
    values: dict[str, int] = {}
    for line in Path("/proc/meminfo").read_text(encoding="utf-8").splitlines():
        key, raw = line.split(":", maxsplit=1)
        amount = int(raw.strip().split()[0])
        values[key] = amount * 1024
    return values


def _owned_processes(
    *,
    pids: set[int],
    process_groups: set[int],
) -> list[dict[str, int | float | str]]:
    result = _run(
        [
            "ps",
            "-eo",
            "pid=,ppid=,pgid=,rss=,pcpu=,comm=",
        ],
        timeout=5,
    )
    processes: list[dict[str, int | float | str]] = []
    for line in result.stdout.splitlines():
        fields = line.split(maxsplit=5)
        if len(fields) != 6:
            continue
        pid, ppid, pgid, rss_kib, cpu_percent, command = fields
        pid_value = int(pid)
        pgid_value = int(pgid)
        if pid_value not in pids and pgid_value not in process_groups:
            continue
        processes.append(
            {
                "pid": pid_value,
                "ppid": int(ppid),
                "pgid": pgid_value,
                "rss_bytes": int(rss_kib) * 1024,
                "cpu_percent": float(cpu_percent),
                "command": command,
            }
        )
    return processes


def _gpu_sample() -> dict[str, int] | None:
    if shutil.which("nvidia-smi") is None:
        return None
    result = _run(
        [
            "nvidia-smi",
            "--query-gpu=memory.used,utilization.gpu",
            "--format=csv,noheader,nounits",
        ],
        timeout=5,
    )
    first_line = result.stdout.strip().splitlines()[0]
    memory_mib, utilization_percent = (int(value.strip()) for value in first_line.split(","))
    return {
        "global_memory_used_mib": memory_mib,
        "global_utilization_percent": utilization_percent,
    }


class ResourceSampler:
    def __init__(
        self,
        *,
        owned_pids: set[int],
        owned_process_groups: set[int],
        storage_root: Path,
        interval_seconds: float,
    ) -> None:
        self._owned_pids = owned_pids
        self._owned_process_groups = owned_process_groups
        self._storage_root = storage_root
        self._interval_seconds = interval_seconds
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, name="s0-resource-sampler")
        self.samples: list[dict[str, Any]] = []
        self.error: str | None = None

    def start(self) -> None:
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        self._thread.join(timeout=max(5.0, self._interval_seconds * 4))
        if self._thread.is_alive():
            raise ProbeError("resource sampler did not stop")
        if self.error is not None:
            raise ProbeError(f"resource sampler failed: {self.error}")

    def _run(self) -> None:
        try:
            while True:
                processes = _owned_processes(
                    pids=self._owned_pids,
                    process_groups=self._owned_process_groups,
                )
                memory = _meminfo()
                disk = shutil.disk_usage(self._storage_root)
                self.samples.append(
                    {
                        "monotonic_ns": time.monotonic_ns(),
                        "owned_rss_bytes": sum(int(process["rss_bytes"]) for process in processes),
                        "owned_cpu_percent": sum(
                            float(process["cpu_percent"]) for process in processes
                        ),
                        "memory_available_bytes": memory["MemAvailable"],
                        "swap_used_bytes": memory["SwapTotal"] - memory["SwapFree"],
                        "disk_free_bytes": disk.free,
                        "gpu": _gpu_sample(),
                        "processes": processes,
                    }
                )
                if self._stop.wait(self._interval_seconds):
                    break
        except Exception as error:  # noqa: BLE001 - preserve factual sampler failure
            self.error = f"{type(error).__name__}: {error}"


def _load_contract(profile_path: Path, speed_factor: int) -> dict[str, int | str]:
    profile = yaml.safe_load(profile_path.read_text(encoding="utf-8"))
    if not isinstance(profile, dict):
        raise ProbeError("qualification profile must be a mapping")
    runtime = profile.get("runtime_acceptance")
    storage = profile.get("storage")
    if not isinstance(runtime, dict) or not isinstance(storage, dict):
        raise ProbeError("qualification profile omitted runtime_acceptance or storage")
    factors = runtime.get("speed_factors")
    if not isinstance(factors, list) or speed_factor not in factors:
        raise ProbeError(f"speed factor {speed_factor} is not declared by the profile")
    return {
        "world": str(runtime["world"]),
        "model": str(runtime["model"]),
        "physics_step_ns": int(runtime["physics_step_ns"]),
        "pause_max_advance_ns": int(runtime["pause_max_advance_ns"]),
        "rtf_tolerance_percent": int(runtime["rtf_tolerance_percent"]),
        "measurement_window_seconds": int(runtime["measurement_window_seconds"]),
        "resource_sample_interval_milliseconds": int(
            runtime["resource_sample_interval_milliseconds"]
        ),
        "max_owned_rss_mib": int(runtime["max_owned_rss_mib"]),
        "max_owned_cpu_percent": int(runtime["max_owned_cpu_percent"]),
        "minimum_available_memory_gib": int(runtime["minimum_available_memory_gib"]),
        "stop_below_gib": int(storage["stop_below_gib"]),
    }


def _resource_report(
    *,
    samples: list[dict[str, Any]],
    contract: dict[str, int | str],
    speed_factor: int,
) -> dict[str, Any]:
    if not samples:
        raise ProbeError("resource sampler produced no samples")
    peak_rss = max(int(sample["owned_rss_bytes"]) for sample in samples)
    peak_cpu = max(float(sample["owned_cpu_percent"]) for sample in samples)
    minimum_memory = min(int(sample["memory_available_bytes"]) for sample in samples)
    minimum_disk = min(int(sample["disk_free_bytes"]) for sample in samples)
    maximum_swap = max(int(sample["swap_used_bytes"]) for sample in samples)
    gpu_samples = [sample["gpu"] for sample in samples if sample["gpu"] is not None]
    checks = {
        "sample_count": len(samples) >= 6,
        "owned_rss": peak_rss <= int(contract["max_owned_rss_mib"]) * MIB,
        "owned_cpu": peak_cpu <= int(contract["max_owned_cpu_percent"]),
        "available_memory": (minimum_memory >= int(contract["minimum_available_memory_gib"]) * GIB),
        "disk_floor": minimum_disk >= int(contract["stop_below_gib"]) * GIB,
    }
    summary: dict[str, Any] = {
        "sample_count": len(samples),
        "peak_owned_rss_bytes": peak_rss,
        "peak_owned_cpu_percent": peak_cpu,
        "minimum_memory_available_bytes": minimum_memory,
        "minimum_disk_free_bytes": minimum_disk,
        "maximum_swap_used_bytes": maximum_swap,
        "gpu_metrics_are_global_and_include_co_tenants": True,
    }
    if gpu_samples:
        summary["peak_global_gpu_memory_used_mib"] = max(
            int(sample["global_memory_used_mib"]) for sample in gpu_samples
        )
        summary["peak_global_gpu_utilization_percent"] = max(
            int(sample["global_utilization_percent"]) for sample in gpu_samples
        )
    return {
        "schema_version": RESOURCE_SCHEMA,
        "speed_factor": speed_factor,
        "actuator_authority": False,
        "checks": checks,
        "summary": summary,
        "samples": samples,
        "verdict": "pass" if all(checks.values()) else "fail",
    }


def _write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--profile", type=Path, required=True)
    parser.add_argument("--speed-factor", type=int, required=True)
    parser.add_argument("--time-output", type=Path, required=True)
    parser.add_argument("--resource-output", type=Path, required=True)
    parser.add_argument("--storage-root", type=Path, required=True)
    parser.add_argument("--owned-pid", action="append", type=int, default=[])
    parser.add_argument("--owned-pgid", action="append", type=int, default=[])
    return parser.parse_args()


def main() -> int:
    args = _arguments()
    contract = _load_contract(args.profile, args.speed_factor)
    sampler = ResourceSampler(
        owned_pids=set(args.owned_pid),
        owned_process_groups=set(args.owned_pgid),
        storage_root=args.storage_root,
        interval_seconds=int(contract["resource_sample_interval_milliseconds"]) / 1000,
    )
    time_report: dict[str, Any] = {
        "schema_version": TIME_SCHEMA,
        "speed_factor": args.speed_factor,
        "world": contract["world"],
        "model": contract["model"],
        "actuator_authority": False,
        "contract": contract,
    }
    exit_code = 1
    sampler.start()
    try:
        world = str(contract["world"])
        _control(world, "pause: true")
        paused_start = _wait_for_pause_state(True)
        time.sleep(1.0)
        paused_end = _stats_snapshot()
        pause_delta = paused_end.sim_time_ns - paused_start.sim_time_ns
        pause_pass = paused_end.paused and 0 <= pause_delta <= int(contract["pause_max_advance_ns"])

        step_before = paused_end
        _control(world, "pause: true, multi_step: 1")
        step_after = _wait_for_iteration(step_before.iterations + 1)
        step_delta = step_after.sim_time_ns - step_before.sim_time_ns
        iteration_delta = step_after.iterations - step_before.iterations
        step_pass = (
            step_after.paused
            and step_delta == int(contract["physics_step_ns"])
            and iteration_delta == 1
        )

        _control(world, "pause: false")
        _wait_for_pause_state(False)
        time.sleep(0.75)
        speed_start = _stats_snapshot()
        time.sleep(int(contract["measurement_window_seconds"]))
        speed_end = _stats_snapshot()
        sim_delta = speed_end.sim_time_ns - speed_start.sim_time_ns
        real_delta = speed_end.real_time_ns - speed_start.real_time_ns
        if real_delta <= 0:
            raise ProbeError("Gazebo real-time delta was not positive")
        measured_rtf = sim_delta / real_delta
        tolerance = int(contract["rtf_tolerance_percent"]) / 100
        lower = args.speed_factor * (1 - tolerance)
        upper = args.speed_factor * (1 + tolerance)
        speed_pass = lower <= measured_rtf <= upper

        time_report.update(
            {
                "pause": {
                    "before": asdict(paused_start),
                    "after": asdict(paused_end),
                    "observed_advance_ns": pause_delta,
                    "pass": pause_pass,
                },
                "single_step": {
                    "before": asdict(step_before),
                    "after": asdict(step_after),
                    "observed_advance_ns": step_delta,
                    "observed_iteration_delta": iteration_delta,
                    "pass": step_pass,
                },
                "speed": {
                    "before": asdict(speed_start),
                    "after": asdict(speed_end),
                    "simulation_advance_ns": sim_delta,
                    "real_advance_ns": real_delta,
                    "measured_rtf": measured_rtf,
                    "accepted_range": [lower, upper],
                    "pass": speed_pass,
                },
            }
        )
        time_report["verdict"] = "pass" if pause_pass and step_pass and speed_pass else "fail"
    except Exception as error:  # noqa: BLE001 - persist the factual probe failure
        time_report["verdict"] = "fail"
        time_report["error"] = f"{type(error).__name__}: {error}"
    finally:
        try:
            _control(str(contract["world"]), "pause: false")
        except Exception as resume_error:  # noqa: BLE001 - report cleanup failure
            time_report["resume_error"] = f"{type(resume_error).__name__}: {resume_error}"
            time_report["verdict"] = "fail"
        try:
            sampler.stop()
            resource_report = _resource_report(
                samples=sampler.samples,
                contract=contract,
                speed_factor=args.speed_factor,
            )
        except Exception as resource_error:  # noqa: BLE001
            resource_report = {
                "schema_version": RESOURCE_SCHEMA,
                "speed_factor": args.speed_factor,
                "actuator_authority": False,
                "verdict": "fail",
                "error": f"{type(resource_error).__name__}: {resource_error}",
                "samples": sampler.samples,
            }
        _write_json(args.time_output, time_report)
        _write_json(args.resource_output, resource_report)

    if time_report["verdict"] == "pass" and resource_report["verdict"] == "pass":
        exit_code = 0
    print(
        json.dumps(
            {
                "time_verdict": time_report["verdict"],
                "resource_verdict": resource_report["verdict"],
                "speed_factor": args.speed_factor,
            },
            sort_keys=True,
        )
    )
    return exit_code


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