#!/usr/bin/env python3
"""Rehearse the public protocol-4 to protocol-5 migration and rollback runbook."""

from __future__ import annotations

import argparse
import fcntl
import hashlib
import json
import os
from pathlib import Path
import shutil
import signal
import subprocess
import sys
import tempfile
import time
from typing import Any, Sequence


class RehearsalError(RuntimeError):
    """A stable migration-rehearsal refusal."""


def _clean_environment() -> dict[str, str]:
    environment = os.environ.copy()
    for name in ("AGCOORD_RUN_ID", "AGCOORD_RUN_KIND", "AGCOORD_STATE_DIR"):
        environment.pop(name, None)
    environment["PYTHONNOUSERSITE"] = "1"
    return environment


def _run(
    arguments: Sequence[str | os.PathLike[str]],
    *,
    cwd: Path | None = None,
    success: bool = True,
) -> subprocess.CompletedProcess[str]:
    command = [os.fspath(argument) for argument in arguments]
    completed = subprocess.run(
        command,
        cwd=cwd,
        env=_clean_environment(),
        stdin=subprocess.DEVNULL,
        text=True,
        capture_output=True,
        check=False,
    )
    if (completed.returncode == 0) != success:
        if completed.returncode == 0:
            raise RehearsalError(f"command unexpectedly succeeded: {command!r}")
        raise RehearsalError(
            f"command exited {completed.returncode}: {command!r}\n"
            f"stdout={completed.stdout}\nstderr={completed.stderr}"
        )
    return completed


def _json_output(completed: subprocess.CompletedProcess[str], subject: str) -> Any:
    try:
        return json.loads(completed.stdout)
    except json.JSONDecodeError as exc:
        raise RehearsalError(f"{subject} did not emit JSON: {completed.stdout!r}") from exc


def _expect_refusal(
    arguments: Sequence[str | os.PathLike[str]],
    *required_fragments: str,
) -> None:
    completed = _run(arguments, success=False)
    report = f"{completed.stdout}\n{completed.stderr}".lower()
    if not all(fragment.lower() in report for fragment in required_fragments):
        raise RehearsalError(
            f"refusal did not contain {required_fragments!r}: {report!r}"
        )


def _owner_fields(state: Path) -> dict[str, str]:
    try:
        raw = (state / "broker.lock").read_text(encoding="utf-8")
    except OSError as exc:
        raise RehearsalError(f"cannot read rehearsal owner metadata: {exc}") from exc
    fields: dict[str, str] = {}
    for line in raw.splitlines():
        if "=" in line:
            name, value = line.split("=", 1)
            fields[name] = value
    return fields


def _lock_is_free(state: Path) -> bool:
    descriptor = os.open(state / "broker.lock", os.O_RDWR)
    try:
        try:
            fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            return False
        fcntl.flock(descriptor, fcntl.LOCK_UN)
        return True
    finally:
        os.close(descriptor)


def _wait_for_owner(state: Path, protocol: int, process: subprocess.Popen[str]) -> None:
    deadline = time.monotonic() + 10
    while time.monotonic() < deadline:
        if process.poll() is not None:
            stdout, stderr = process.communicate()
            raise RehearsalError(
                f"protocol-{protocol} rehearsal owner exited early: "
                f"stdout={stdout!r}, stderr={stderr!r}"
            )
        try:
            fields = _owner_fields(state)
            if fields.get("protocol") == str(protocol) and not _lock_is_free(state):
                return
        except (OSError, RehearsalError):
            pass
        time.sleep(0.02)
    raise RehearsalError(f"protocol-{protocol} rehearsal owner did not start")


def _wait_for_unlock(state: Path) -> None:
    deadline = time.monotonic() + 10
    while time.monotonic() < deadline:
        try:
            if _lock_is_free(state):
                return
        except OSError:
            pass
        time.sleep(0.02)
    raise RehearsalError("rehearsal owner did not release its spool lock")


def _start_python_owner(python: Path, state: Path) -> subprocess.Popen[str]:
    process = subprocess.Popen(
        [
            os.fspath(python),
            "-m",
            "agcoord.queue",
            "serve",
            "--state-dir",
            os.fspath(state),
            "--idle-seconds",
            "300",
        ],
        env=_clean_environment(),
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )
    _wait_for_owner(state, 4, process)
    return process


def _stop_python_owner(process: subprocess.Popen[str], state: Path) -> None:
    if process.poll() is None:
        process.send_signal(signal.SIGTERM)
        try:
            process.communicate(timeout=10)
        except subprocess.TimeoutExpired as exc:
            process.kill()
            process.communicate()
            raise RehearsalError("protocol-4 rehearsal owner did not stop") from exc
    if process.returncode != 0:
        stdout, stderr = process.communicate()
        raise RehearsalError(
            f"protocol-4 rehearsal owner exited {process.returncode}: "
            f"stdout={stdout!r}, stderr={stderr!r}"
        )
    _wait_for_unlock(state)


def _stop_native_owner(state: Path, broker: Path) -> None:
    if not (state / "broker.lock").exists() or _lock_is_free(state):
        return
    fields = _owner_fields(state)
    if fields.get("protocol") != "5" or fields.get("implementation") != "rust-native":
        raise RehearsalError("refusing to stop a non-native rehearsal owner")
    try:
        pid = int(fields["pid"])
    except (KeyError, ValueError) as exc:
        raise RehearsalError("native rehearsal owner has invalid PID metadata") from exc
    executable = Path(f"/proc/{pid}/exe")
    try:
        if not os.path.samefile(executable, broker):
            raise RehearsalError("native rehearsal owner executable identity changed")
        os.kill(pid, signal.SIGTERM)
    except ProcessLookupError:
        pass
    except OSError as exc:
        raise RehearsalError(f"cannot stop native rehearsal owner: {exc}") from exc
    _wait_for_unlock(state)


def _initialize_checkout(checkout: Path) -> None:
    checkout.mkdir()
    _run(["git", "init", "-q"], cwd=checkout)
    _run(["git", "config", "user.name", "AGCoord migration rehearsal"], cwd=checkout)
    _run(
        ["git", "config", "user.email", "migration-rehearsal@example.invalid"],
        cwd=checkout,
    )
    (checkout / "tracked.txt").write_text("migration rehearsal\n", encoding="utf-8")
    _run(["git", "add", "tracked.txt"], cwd=checkout)
    _run(["git", "commit", "-q", "-m", "fixture"], cwd=checkout)


def _write_config(state: Path, broker: Path) -> None:
    state.mkdir(mode=0o700)
    config = {
        "capacities": {"jobs": 1},
        "native_broker": {
            "path": os.fspath(broker),
            "allow_development": True,
        },
    }
    path = state / "config.json"
    path.write_text(json.dumps(config, sort_keys=True), encoding="utf-8")
    path.chmod(0o600)


def _seed_legacy_history(python: Path, state: Path, checkout: Path) -> str:
    program = r'''
import json
import os
import sys
import time
from agcoord.queue import CoordinatorClient

client = CoordinatorClient(state_dir=sys.argv[1], autostart=False)
run_id = client.submit(
    ["/bin/sh", "-c", "printf 'legacy migration rehearsal\\n'"],
    checkout=sys.argv[2],
    label="legacy migration rehearsal",
    resources={"jobs": 1},
    caller_pid=os.getpid(),
)
deadline = time.monotonic() + 10
while time.monotonic() < deadline:
    row = client.status(run_id)
    if row["status"] not in {"queued", "running"}:
        print(json.dumps(row, sort_keys=True))
        raise SystemExit(0 if row["status"] == "passed" else 1)
    time.sleep(0.02)
raise SystemExit("legacy rehearsal run did not finish")
'''
    completed = _run([python, "-c", program, state, checkout])
    row = _json_output(completed, "legacy rehearsal run")
    return str(row["run_id"])


def _assert_python_history(
    python: Path,
    state: Path,
    run_ids: Sequence[str],
) -> None:
    program = r'''
import json
import sys
from agcoord.queue import CoordinatorClient

client = CoordinatorClient(state_dir=sys.argv[1], autostart=False)
print(json.dumps({run_id: client.status(run_id)["status"] for run_id in sys.argv[2:]}))
'''
    completed = _run([python, "-c", program, state, *run_ids])
    statuses = _json_output(completed, "protocol-4 history inspection")
    if statuses != {run_id: "passed" for run_id in run_ids}:
        raise RehearsalError(f"migration history changed: {statuses!r}")


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def rehearse(python: Path, agc: Path, broker: Path) -> dict[str, Any]:
    identity = _json_output(_run([broker, "identity", "--json"]), "native identity")
    if (
        identity.get("protocol") != 5
        or identity.get("implementation") != "rust-native"
    ):
        raise RehearsalError("selected migration executable is not a protocol-5 Rust broker")

    with tempfile.TemporaryDirectory(prefix="agcoord-migration-rehearsal-") as raw:
        root = Path(raw)
        state = root / "state"
        checkout = root / "checkout"
        operator_backup = root / "operator-backup"
        rollback_rehearsal = root / "rollback-rehearsal"
        _initialize_checkout(checkout)
        _write_config(state, broker)
        owners: list[tuple[subprocess.Popen[str], Path]] = []
        try:
            legacy_owner = _start_python_owner(python, state)
            owners.append((legacy_owner, state))
            legacy_run_id = _seed_legacy_history(python, state, checkout)
            _expect_refusal(
                [agc, "--json", "--state-dir", state, "list"],
                "legacy protocol-4",
                "migrate",
            )
            _stop_python_owner(legacy_owner, state)
            owners.remove((legacy_owner, state))

            _expect_refusal(
                [agc, "--json", "--state-dir", state, "list"],
                "protocol 4",
                "migrate",
            )
            shutil.copytree(state, operator_backup, copy_function=shutil.copy2)
            untouched_backup_sha256 = _sha256(operator_backup / "queue.sqlite3")
            shutil.copytree(
                operator_backup,
                rollback_rehearsal,
                copy_function=shutil.copy2,
            )
            migrated_backup = _json_output(
                _run([broker, "migrate", "--state-dir", rollback_rehearsal]),
                "backup migration",
            )
            rolled_back_backup = _json_output(
                _run([broker, "rollback", "--state-dir", rollback_rehearsal]),
                "backup rollback",
            )
            if migrated_backup != {"changed": True, "from_protocol": 4, "to_protocol": 5}:
                raise RehearsalError(f"unexpected backup migration: {migrated_backup!r}")
            if rolled_back_backup != {"changed": True, "from_protocol": 5, "to_protocol": 4}:
                raise RehearsalError(f"unexpected backup rollback: {rolled_back_backup!r}")
            backup_owner = _start_python_owner(python, rollback_rehearsal)
            owners.append((backup_owner, rollback_rehearsal))
            _assert_python_history(
                python,
                rollback_rehearsal,
                [legacy_run_id],
            )
            _stop_python_owner(backup_owner, rollback_rehearsal)
            owners.remove((backup_owner, rollback_rehearsal))

            migrated = _json_output(
                _run([agc, "--json", "--state-dir", state, "migrate"]),
                "live migration",
            )
            if migrated != {"changed": True, "from_protocol": 4, "to_protocol": 5}:
                raise RehearsalError(f"unexpected live migration: {migrated!r}")
            backups = list(state.glob("queue.sqlite3.protocol4*.bak"))
            if not backups or any(path.stat().st_mode & 0o077 for path in backups):
                raise RehearsalError("migration did not retain a private protocol-4 backup")

            native_row = _json_output(
                _run(
                    [
                        agc,
                        "--json",
                        "--state-dir",
                        state,
                        "run",
                        "--checkout",
                        checkout,
                        "--",
                        "/bin/true",
                    ]
                ),
                "native rehearsal run",
            )
            if native_row.get("status") != "passed":
                raise RehearsalError(f"native rehearsal run failed: {native_row!r}")
            native_run_id = str(native_row["run_id"])
            _stop_native_owner(state, broker)

            rolled_back = _json_output(
                _run([broker, "rollback", "--state-dir", state]),
                "live rollback",
            )
            if rolled_back != {"changed": True, "from_protocol": 5, "to_protocol": 4}:
                raise RehearsalError(f"unexpected live rollback: {rolled_back!r}")
            rollback_owner = _start_python_owner(python, state)
            owners.append((rollback_owner, state))
            _assert_python_history(
                python,
                state,
                [legacy_run_id, native_run_id],
            )
            _stop_python_owner(rollback_owner, state)
            owners.remove((rollback_owner, state))

            remigrated = _json_output(
                _run([agc, "--json", "--state-dir", state, "migrate"]),
                "final migration",
            )
            if remigrated != {"changed": True, "from_protocol": 4, "to_protocol": 5}:
                raise RehearsalError(f"unexpected final migration: {remigrated!r}")
            _expect_refusal(
                [
                    python,
                    "-c",
                    "import sys; from agcoord.queue import CoordinatorBroker; "
                    "CoordinatorBroker(sys.argv[1], idle_timeout=None)",
                    state,
                ],
                "protocol",
                "need 4",
            )
            final_snapshot = _json_output(
                _run([agc, "--json", "--state-dir", state, "list"]),
                "final protocol-5 snapshot",
            )
            if final_snapshot.get("protocol") != 5:
                raise RehearsalError("final client did not join a protocol-5 owner")
            _stop_native_owner(state, broker)
            if _sha256(operator_backup / "queue.sqlite3") != untouched_backup_sha256:
                raise RehearsalError("operator backup changed during the rehearsal")
            return {
                "backup_sha256": untouched_backup_sha256,
                "broker_build": identity["build"],
                "broker_version": identity["version"],
                "final_protocol": 5,
                "legacy_run_id": legacy_run_id,
                "native_run_id": native_run_id,
                "rollback_protocol": 4,
            }
        finally:
            for process, owned_state in owners:
                if process.poll() is None:
                    process.terminate()
                    try:
                        process.communicate(timeout=5)
                    except subprocess.TimeoutExpired:
                        process.kill()
                        process.communicate()
                try:
                    _wait_for_unlock(owned_state)
                except (OSError, RehearsalError):
                    pass
            try:
                _stop_native_owner(state, broker)
            except (OSError, RehearsalError):
                pass


def main() -> int:
    parser = argparse.ArgumentParser(
        description="rehearse native migration and rollback in owned temporary state"
    )
    parser.add_argument("--python", type=Path, default=Path(sys.executable))
    parser.add_argument("--agc", type=Path, required=True)
    parser.add_argument("--broker", type=Path, required=True)
    arguments = parser.parse_args()
    try:
        receipt = rehearse(
            Path(os.path.abspath(arguments.python.expanduser())),
            arguments.agc.expanduser().resolve(),
            arguments.broker.expanduser().resolve(),
        )
    except (OSError, RehearsalError) as exc:
        print(f"native migration rehearsal refused: {exc}", file=sys.stderr)
        return 1
    print(json.dumps(receipt, indent=2, sort_keys=True))
    return 0


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