"""``bench`` sub-commands.

Phase C ports the manual PowerShell benchmark/burn-in tooling to Python so
the Linux host no longer needs ``pwsh``:

* ``bench run``     — concurrency burn-in (replaces ``Test-CapacityBurnIn.ps1``
  + ``Start-BurnInTest*.ps1``).
* ``bench measure`` — phase timing (replaces ``Measure-CIBenchmark.ps1``).

Design rule: import ``backends.load_backend`` and the existing transports
only — never ``WorkstationVmrunBackend`` directly (keeps the Phase D ESXi
path open). No ``Start-Job``/``pwsh``/``$IsWindows``.
"""

from __future__ import annotations

import contextlib
import json
import shutil
import socket
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING

import click

from ci_orchestrator.backends import load_backend
from ci_orchestrator.backends.errors import BackendError, BackendNotAvailable
from ci_orchestrator.backends.protocol import VmHandle
from ci_orchestrator.config import Config, load_config

if TYPE_CHECKING:  # pragma: no cover
    from ci_orchestrator.backends.protocol import VmBackend


# Clone directories left behind by a job are named ``Clone_<job_id>_...``
# (see commands/job.py:_make_clone_name and commands/vm.py:vm_new).
_CLONE_PREFIX = "Clone_"
# A vm-start.lock older than this (seconds) is considered stale leftover.
_STALE_LOCK_SECONDS = 5 * 60
# Linux guests probe SSH/22; everything else probes WinRM/5986.
_SSH_PORT = 22
_WINRM_PORT = 5986


def _now() -> float:
    return time.monotonic()


def _timestamp() -> str:
    return datetime.now(tz=UTC).strftime("%Y%m%d-%H%M%S")


# ─────────────────────────────────────────────────────────────────── group


@click.group("bench")
def bench() -> None:
    """Capacity burn-in and phase-timing benchmarks."""


# ═══════════════════════════════════════════════════════════════ bench run


@dataclass
class JobResult:
    """Outcome of a single burn-in job slot."""

    job_id: str
    slot: int
    exit_code: int
    status: str
    elapsed_sec: float
    error: str = ""


@dataclass
class RoundResult:
    """Outcome of one burn-in round (all slots + cleanup assertions)."""

    round: int
    passed: int
    failed: int
    status: str
    elapsed_sec: float
    orphans: list[str] = field(default_factory=list)
    stale_lock: bool = False
    jobs: list[JobResult] = field(default_factory=list)


def _run_one_job(
    *,
    job_id: str,
    guest_os: str,
    build_command: str,
    clone_base_dir: Path,
) -> tuple[int, str]:
    """Run a single job end to end, returning ``(exit_code, error_message)``.

    Invokes the in-process ``commands.job.job`` command (the same code path
    production uses) rather than re-spawning ``python -m ci_orchestrator`` or
    a PowerShell child — closes the ``Start-Job``/``$IsWindows`` shim bug.

    Isolated in its own module-level function so the burn-in harness can be
    unit tested by monkeypatching it (no real VMs needed).
    """
    from click.testing import CliRunner  # local import: drives the job command

    from ci_orchestrator.commands.job import job

    args = [
        "--job-id",
        job_id,
        "--repo-url",
        "",
        "--branch",
        "main",
        "--guest-os",
        guest_os,
        "--clone-base-dir",
        str(clone_base_dir),
        "--skip-artifact",
    ]
    if build_command:
        args += ["--build-command", build_command]
    result = CliRunner().invoke(job, args, catch_exceptions=True)
    if result.exit_code == 0:
        return 0, ""
    tail = result.output.strip().splitlines()
    return result.exit_code, tail[-1] if tail else ""


def _find_orphans(clone_base: Path, job_ids: list[str]) -> list[Path]:
    """Return clone dirs in ``clone_base`` belonging to the given jobs."""
    if not clone_base.is_dir():
        return []
    prefixes = tuple(f"{_CLONE_PREFIX}{jid}_" for jid in job_ids)
    return [
        entry
        for entry in clone_base.iterdir()
        if entry.is_dir() and entry.name.startswith(prefixes)
    ]


def _stale_lock(lock_file: Path, max_age_seconds: int = _STALE_LOCK_SECONDS) -> bool:
    """True if ``lock_file`` exists and is older than ``max_age_seconds``."""
    if not lock_file.is_file():
        return False
    try:
        mtime = lock_file.stat().st_mtime
    except OSError:
        return False
    return (time.time() - mtime) > max_age_seconds


def _run_round(
    *,
    round_no: int,
    concurrency: int,
    guest_os: str,
    build_command: str,
    clone_base_dir: Path,
    lock_file: Path,
) -> RoundResult:
    """Launch ``concurrency`` concurrent jobs and assert cleanup invariants."""
    click.echo(f"\n[bench run] Round {round_no}: launching {concurrency} job(s)...")
    started = _now()
    ts = _timestamp()
    slots = {
        slot: f"burnin-r{round_no}-j{slot}-{ts}"
        for slot in range(1, concurrency + 1)
    }

    jobs: list[JobResult] = []
    with ThreadPoolExecutor(max_workers=concurrency) as pool:
        futures = {
            pool.submit(
                _run_one_job,
                job_id=job_id,
                guest_os=guest_os,
                build_command=build_command,
                clone_base_dir=clone_base_dir,
            ): (slot, job_id)
            for slot, job_id in slots.items()
        }
        for fut in as_completed(futures):
            slot, job_id = futures[fut]
            try:
                exit_code, error = fut.result()
            except Exception as exc:  # pragma: no cover - defensive
                exit_code, error = 1, str(exc)
            status = "PASS" if exit_code == 0 else "FAIL"
            jobs.append(
                JobResult(
                    job_id=job_id,
                    slot=slot,
                    exit_code=exit_code,
                    status=status,
                    elapsed_sec=round(_now() - started, 2),
                    error=error,
                )
            )

    jobs.sort(key=lambda j: j.slot)
    passed = sum(1 for j in jobs if j.status == "PASS")
    failed = len(jobs) - passed

    # ── Assertion: no orphaned clone directories ─────────────────────────
    orphans = _find_orphans(clone_base_dir, list(slots.values()))
    if orphans:
        failed += 1

    # ── Assertion: no stale vm-start.lock ────────────────────────────────
    stale = _stale_lock(lock_file)
    if stale:
        failed += 1

    status = "PASS" if failed == 0 else "FAIL"
    result = RoundResult(
        round=round_no,
        passed=passed,
        failed=failed,
        status=status,
        elapsed_sec=round(_now() - started, 2),
        orphans=[str(p) for p in orphans],
        stale_lock=stale,
        jobs=jobs,
    )

    for j in jobs:
        line = f"  slot {j.slot}  {j.job_id:<32} {j.status:<5} exit={j.exit_code}"
        if j.error:
            line += f"  ({j.error})"
        click.echo(line)
    click.echo(
        f"  clone cleanup : {'OK' if not orphans else f'FAIL ({len(orphans)} orphan)'}"
    )
    for p in result.orphans:
        click.echo(f"    orphan: {p}")
    click.echo(f"  lock cleanup  : {'OK' if not stale else 'FAIL (stale lock)'}")
    click.echo(
        f"  [Round {round_no}] {status} "
        f"({passed} PASS, {failed} FAIL, {result.elapsed_sec}s)"
    )
    return result


def _write_run_report(report: dict[str, object], json_out: Path) -> None:
    json_out.parent.mkdir(parents=True, exist_ok=True)
    json_out.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")


@bench.command("run")
@click.option(
    "--guest-os",
    type=click.Choice(["windows", "linux"], case_sensitive=False),
    default="linux",
    show_default=True,
    help="Guest OS family of the build VMs under test.",
)
@click.option(
    "--concurrency",
    type=click.IntRange(1, 32),
    default=4,
    show_default=True,
    help="Number of concurrent jobs per round.",
)
@click.option(
    "--rounds",
    type=click.IntRange(1, 1000),
    default=10,
    show_default=True,
    help="Number of sequential rounds to run.",
)
@click.option(
    "--build-command",
    "build_command",
    default="",
    help="Build command each job runs (default: a no-op sleep/ping).",
)
@click.option(
    "--clone-base-dir",
    "clone_base_dir",
    default=None,
    help="Override clone base dir (defaults to config.paths.build_vms).",
)
@click.option(
    "--json-out",
    "json_out",
    default=None,
    help="Path for the JSON report (default: CI_ARTIFACTS/bench/<ts>.json).",
)
def bench_run(
    guest_os: str,
    concurrency: int,
    rounds: int,
    build_command: str,
    clone_base_dir: str | None,
    json_out: str | None,
) -> None:
    """Run a concurrency burn-in and assert cleanup invariants per round.

    Replaces ``Test-CapacityBurnIn.ps1`` + ``Start-BurnInTest*.ps1``. For
    each of ``--rounds`` rounds it launches ``--concurrency`` concurrent
    jobs (in-process, via :func:`concurrent.futures`), then asserts: every
    job exited 0, no orphaned clone dir remains in the clone base dir, and
    no ``vm-start.lock`` older than 5 minutes is left behind.
    """
    guest_os = guest_os.lower()
    config = load_config()
    base = Path(clone_base_dir) if clone_base_dir else config.paths.build_vms
    lock_file = base / "vm-start.lock"

    started_at = datetime.now(tz=UTC)
    out_path = (
        Path(json_out)
        if json_out
        else config.paths.artifacts / "bench" / f"{_timestamp()}.json"
    )

    click.echo("=" * 60)
    click.echo("[bench run] capacity burn-in")
    click.echo(f"  guest-os    : {guest_os}")
    click.echo(f"  concurrency : {concurrency}")
    click.echo(f"  rounds      : {rounds}")
    click.echo(f"  clone base  : {base}")
    click.echo(f"  build cmd   : {build_command or '(default no-op)'}")
    click.echo("=" * 60)

    round_results: list[RoundResult] = []
    for r in range(1, rounds + 1):
        round_results.append(
            _run_round(
                round_no=r,
                concurrency=concurrency,
                guest_os=guest_os,
                build_command=build_command,
                clone_base_dir=base,
                lock_file=lock_file,
            )
        )

    total_pass = sum(rr.passed for rr in round_results)
    total_fail = sum(rr.failed for rr in round_results)
    rounds_failed = sum(1 for rr in round_results if rr.status == "FAIL")
    overall = "PASS" if total_fail == 0 else "FAIL"

    # ── Summary table ────────────────────────────────────────────────────
    click.echo("\n" + "=" * 60)
    click.echo("[bench run] SUMMARY")
    click.echo(f"  {'Round':<7}{'Pass':<6}{'Fail':<6}{'Elapsed':<10}Status")
    click.echo(f"  {'-' * 5:<7}{'-' * 4:<6}{'-' * 4:<6}{'-' * 8:<10}{'-' * 6}")
    for rr in round_results:
        click.echo(
            f"  {rr.round:<7}{rr.passed:<6}{rr.failed:<6}"
            f"{str(rr.elapsed_sec) + 's':<10}{rr.status}"
        )
    click.echo("")
    click.echo(f"  Total jobs : {rounds * concurrency}")
    click.echo(f"  Passed     : {total_pass}")
    click.echo(f"  Failed     : {total_fail}")
    click.echo(f"  Rounds OK  : {rounds - rounds_failed} / {rounds}")
    click.echo(f"\n  OVERALL: {overall}")
    click.echo("=" * 60)

    report: dict[str, object] = {
        "ts": started_at.isoformat(),
        "guestOS": guest_os,
        "concurrency": concurrency,
        "rounds": rounds,
        "buildCommand": build_command,
        "totalJobs": rounds * concurrency,
        "totalPass": total_pass,
        "totalFail": total_fail,
        "roundsFailed": rounds_failed,
        "overall": overall,
        "results": [asdict(rr) for rr in round_results],
    }
    try:
        _write_run_report(report, out_path)
        click.echo(f"[bench run] report written: {out_path}")
    except OSError as exc:
        click.echo(f"[bench run] could not write report: {exc}", err=True)

    if total_fail > 0:
        raise SystemExit(1)


# ═══════════════════════════════════════════════════════════ bench measure


@dataclass
class PhaseTiming:
    """Per-iteration phase timings, mirroring ``Measure-CIBenchmark.ps1``.

    Field names of the serialised record (see :meth:`to_record`) are kept
    identical to the PowerShell version for trend-log continuity.
    """

    iteration: int
    run_id: str
    template: str
    snapshot: str
    guest_os: str
    clone_sec: float | None = None
    start_sec: float | None = None
    ip_sec: float | None = None
    ready_sec: float | None = None
    destroy_sec: float | None = None
    delta_kb: int | None = None
    vm_ip: str | None = None
    error: str | None = None

    @property
    def total_boot_sec(self) -> float | None:
        if self.ip_sec is None or self.ready_sec is None:
            return None
        return round(
            (self.clone_sec or 0.0)
            + (self.start_sec or 0.0)
            + self.ip_sec
            + self.ready_sec,
            2,
        )

    def to_record(self, ts: str) -> dict[str, object]:
        """Build a JSON record with the legacy ``Measure-CIBenchmark`` keys."""
        return {
            "ts": ts,
            "runId": self.run_id,
            "iteration": self.iteration,
            "template": self.template,
            "snapshot": self.snapshot,
            "guestOS": self.guest_os,
            "cloneSec": self.clone_sec,
            "startSec": self.start_sec,
            "ipSec": self.ip_sec,
            "readySec": self.ready_sec,
            "destroySec": self.destroy_sec,
            "totalBootSec": self.total_boot_sec,
            "deltaKB": self.delta_kb,
            "vmIP": self.vm_ip,
            "error": self.error,
        }


def _dir_size_kb(path: Path) -> int | None:
    """Total size of files under ``path`` in KiB, or None on error."""
    try:
        total = sum(p.stat().st_size for p in path.rglob("*") if p.is_file())
    except OSError:
        return None
    return round(total / 1024)


def _tcp_open(host: str, port: int, timeout: float = 3.0) -> bool:
    """Return True if a TCP connection to ``host:port`` succeeds."""
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except OSError:
        return False


def _wait_port(host: str, port: int, deadline: float, poll: float = 3.0) -> bool:
    while _now() < deadline:
        if _tcp_open(host, port):
            return True
        time.sleep(poll)
    return False


def _measure_iteration(
    *,
    backend: VmBackend,
    iteration: int,
    template: Path,
    snapshot: str,
    guest_os: str,
    clone_base_dir: Path,
    transport_port: int,
    timeout_seconds: float,
) -> PhaseTiming:
    """Clone → start → IP → transport-ready → destroy, timing each phase."""
    run_id = f"bench-{_timestamp()}-{iteration}"
    clone_name = f"{_CLONE_PREFIX}{run_id}"
    clone_dir = clone_base_dir / clone_name
    clone_vmx = clone_dir / f"{clone_name}.vmx"
    timing = PhaseTiming(
        iteration=iteration,
        run_id=run_id,
        template=template.name,
        snapshot=snapshot,
        guest_os=guest_os,
    )
    handle: VmHandle | None = None
    try:
        # ── clone ────────────────────────────────────────────────────────
        click.echo(f"[bench measure] iteration {iteration}: cloning...")
        t0 = _now()
        handle = backend.clone_linked(
            template=str(template),
            snapshot=snapshot,
            name=clone_name,
            destination=str(clone_vmx),
        )
        timing.clone_sec = round(_now() - t0, 2)
        timing.delta_kb = _dir_size_kb(clone_dir)

        # ── start ────────────────────────────────────────────────────────
        click.echo("[bench measure] starting...")
        t0 = _now()
        backend.start(handle, headless=True)
        timing.start_sec = round(_now() - t0, 2)

        # ── IP acquire ────────────────────────────────────────────────────
        click.echo("[bench measure] waiting for guest IP...")
        t0 = _now()
        deadline = t0 + timeout_seconds
        vm_ip: str | None = None
        while _now() < deadline:
            vm_ip = backend.get_ip(handle, timeout=min(5.0, deadline - _now()))
            if vm_ip:
                break
        if not vm_ip:
            raise TimeoutError(f"no guest IP after {timeout_seconds}s")
        timing.ip_sec = round(_now() - t0, 2)
        timing.vm_ip = vm_ip

        # ── transport ready ───────────────────────────────────────────────
        click.echo(f"[bench measure] waiting for port {transport_port}...")
        t0 = _now()
        if not _wait_port(vm_ip, transport_port, _now() + timeout_seconds):
            raise TimeoutError(
                f"port {transport_port} on {vm_ip} not open after {timeout_seconds}s"
            )
        timing.ready_sec = round(_now() - t0, 2)
    except (BackendError, TimeoutError, OSError) as exc:
        timing.error = str(exc)
        click.echo(f"[bench measure] iteration {iteration} failed: {exc}", err=True)
    finally:
        # ── destroy ───────────────────────────────────────────────────────
        t0 = _now()
        if handle is not None:
            with contextlib.suppress(BackendError):
                backend.stop(handle, hard=True)
            with contextlib.suppress(BackendError):
                backend.delete(handle)
        if clone_dir.exists():
            shutil.rmtree(clone_dir, ignore_errors=True)
        timing.destroy_sec = round(_now() - t0, 2)
    return timing


def _resolve_template(config: Config, guest_os: str) -> tuple[Path, str]:
    """Pick a default template VMX + snapshot for ``guest_os``."""
    if guest_os == "linux":
        return (
            config.paths.templates / "LinuxBuild2404" / "LinuxBuild2404.vmx",
            "BaseClean-Linux",
        )
    return (
        config.paths.templates / "WinBuild2025" / "WinBuild2025.vmx",
        "BaseClean",
    )


@bench.command("measure")
@click.option(
    "--guest-os",
    type=click.Choice(["windows", "linux"], case_sensitive=False),
    default="linux",
    show_default=True,
    help="Guest OS family of the VM to benchmark.",
)
@click.option(
    "--iterations",
    type=click.IntRange(1, 100),
    default=1,
    show_default=True,
    help="Number of ephemeral VMs to time.",
)
@click.option(
    "--template-path",
    "template_path",
    default=None,
    help="Template VMX path (defaults to the per-guest-os template).",
)
@click.option(
    "--snapshot-name",
    "snapshot_name",
    default=None,
    help="Snapshot to clone from (defaults to BaseClean / BaseClean-Linux).",
)
@click.option(
    "--clone-base-dir",
    "clone_base_dir",
    default=None,
    help="Override clone base dir (defaults to config.paths.build_vms).",
)
@click.option(
    "--timeout-seconds",
    "timeout_seconds",
    type=click.FloatRange(30.0, 1200.0),
    default=300.0,
    show_default=True,
    help="Per-iteration timeout for IP and transport readiness.",
)
@click.option(
    "--json-out",
    "json_out",
    default=None,
    help="Path to append JSONL records (default: CI_ARTIFACTS/benchmark.jsonl).",
)
def bench_measure(
    guest_os: str,
    iterations: int,
    template_path: str | None,
    snapshot_name: str | None,
    clone_base_dir: str | None,
    timeout_seconds: float,
    json_out: str | None,
) -> None:
    """Measure clone/start/IP/transport/destroy phase timings.

    Replaces ``Measure-CIBenchmark.ps1``. Records are appended to
    ``benchmark.jsonl`` with the same field names so historical trend data
    stays comparable.
    """
    guest_os = guest_os.lower()
    config = load_config()

    default_tpl, default_snap = _resolve_template(config, guest_os)
    template = Path(template_path) if template_path else default_tpl
    snapshot = snapshot_name or default_snap
    base = Path(clone_base_dir) if clone_base_dir else config.paths.build_vms
    transport_port = _SSH_PORT if guest_os == "linux" else _WINRM_PORT
    out_path = Path(json_out) if json_out else config.paths.artifacts / "benchmark.jsonl"

    if not template.is_file():
        raise click.ClickException(f"Template VMX not found: {template}")

    try:
        backend = load_backend(config)
    except BackendNotAvailable as exc:
        raise click.ClickException(f"backend unavailable: {exc}") from exc

    base.mkdir(parents=True, exist_ok=True)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    click.echo("=" * 60)
    click.echo("[bench measure] phase timing")
    click.echo(f"  template   : {template.name}")
    click.echo(f"  snapshot   : {snapshot}")
    click.echo(f"  guest-os   : {guest_os} (port {transport_port})")
    click.echo(f"  iterations : {iterations}")
    click.echo("=" * 60)

    timings: list[PhaseTiming] = []
    for i in range(1, iterations + 1):
        timing = _measure_iteration(
            backend=backend,
            iteration=i,
            template=template,
            snapshot=snapshot,
            guest_os=guest_os,
            clone_base_dir=base,
            transport_port=transport_port,
            timeout_seconds=timeout_seconds,
        )
        timings.append(timing)
        record = timing.to_record(datetime.now(tz=UTC).isoformat())
        try:
            with out_path.open("a", encoding="utf-8") as fh:
                fh.write(json.dumps(record) + "\n")
        except OSError as exc:
            click.echo(f"[bench measure] could not append jsonl: {exc}", err=True)

    # ── Summary table ────────────────────────────────────────────────────
    click.echo("\n[bench measure] Results")
    header = f"  {'Iter':<5}{'Clone':<8}{'Start':<8}{'IP':<8}{'Ready':<8}{'Destroy':<9}{'Boot':<8}IP"
    click.echo(header)
    for t in timings:
        click.echo(
            f"  {t.iteration:<5}"
            f"{t.clone_sec!s:<8}{t.start_sec!s:<8}{t.ip_sec!s:<8}"
            f"{t.ready_sec!s:<8}{t.destroy_sec!s:<9}"
            f"{t.total_boot_sec!s:<8}{t.vm_ip or '-'}"
        )
        if t.error:
            click.echo(f"        error: {t.error}")

    successful = [t for t in timings if t.error is None and t.total_boot_sec is not None]
    if len(successful) > 1:
        boots = [t.total_boot_sec for t in successful if t.total_boot_sec is not None]
        avg = round(sum(boots) / len(boots), 2)
        click.echo(f"\n  Average boot-to-ready ({len(successful)} ok): {avg}s")
    click.echo(f"\n[bench measure] records appended to: {out_path}")


__all__ = ["bench"]