#!/usr/bin/env python3
"""Exercise GovOPlaN worker delivery guarantees against a real Redis broker."""

from __future__ import annotations

import argparse
from datetime import UTC, datetime
import json
import os
from pathlib import Path
import signal
import subprocess
import sys
import tempfile
import time
from typing import Any
from urllib.parse import urlsplit, urlunsplit
import uuid


TERMINAL_STATES = {"FAILURE", "REVOKED", "SUCCESS"}


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Prove Celery publish/consume, retry, warm shutdown, and worker-loss "
            "redelivery against an isolated Redis database."
        )
    )
    parser.add_argument(
        "--redis-url",
        default=os.environ.get("GOVOPLAN_WORKER_DRILL_REDIS_URL", ""),
    )
    parser.add_argument("--python", default=sys.executable)
    parser.add_argument("--timeout-seconds", type=float, default=120.0)
    parser.add_argument(
        "--visibility-timeout-seconds",
        type=int,
        default=30,
    )
    parser.add_argument("--output", type=Path)
    return parser


def _redacted_redis_url(value: str) -> str:
    parsed = urlsplit(value)
    hostname = parsed.hostname or ""
    port = f":{parsed.port}" if parsed.port else ""
    return urlunsplit((parsed.scheme, f"{hostname}{port}", parsed.path, "", ""))


def _wait_until(predicate, *, timeout_seconds: float, detail: str):
    deadline = time.monotonic() + timeout_seconds
    last_value: Any = None
    while time.monotonic() < deadline:
        last_value = predicate()
        if last_value:
            return last_value
        time.sleep(0.2)
    raise TimeoutError(f"Timed out waiting for {detail}; last value: {last_value!r}")


def _worker_command(python: str, hostname: str) -> list[str]:
    return [
        python,
        "-m",
        "celery",
        "-A",
        "govoplan_core.celery_app:celery",
        "worker",
        "--pool",
        "solo",
        "--queues",
        "default",
        "--hostname",
        hostname,
        "--loglevel",
        "WARNING",
        "--without-mingle",
    ]


def _start_worker(
    *,
    python: str,
    hostname: str,
    environment: dict[str, str],
    log_path: Path,
) -> tuple[subprocess.Popen[bytes], Any]:
    log = log_path.open("ab", buffering=0)
    process = subprocess.Popen(
        _worker_command(python, hostname),
        env=environment,
        stdin=subprocess.DEVNULL,
        stdout=log,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    return process, log


def _stop_worker(
    process: subprocess.Popen[bytes] | None,
    log: Any,
    *,
    timeout_seconds: float = 30.0,
) -> None:
    if process is not None and process.poll() is None:
        process.send_signal(signal.SIGTERM)
        try:
            process.wait(timeout=timeout_seconds)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait(timeout=10)
    if log is not None:
        log.close()


def _wait_for_worker(app, process: subprocess.Popen[bytes], prefix: str, timeout: float) -> str:
    def ping() -> str | None:
        if process.poll() is not None:
            raise RuntimeError(f"Worker exited before readiness with code {process.returncode}")
        replies = app.control.ping(timeout=1.0) or []
        for reply in replies:
            for hostname, payload in reply.items():
                if hostname.startswith(prefix) and payload.get("ok") == "pong":
                    return hostname
        return None

    return _wait_until(ping, timeout_seconds=timeout, detail=f"worker {prefix}")


def _wait_for_started(result, *, timeout_seconds: float) -> dict[str, Any]:
    def started() -> dict[str, Any] | None:
        state = result.state
        if state in TERMINAL_STATES and state != "SUCCESS":
            raise RuntimeError(f"Probe {result.id} became {state}: {result.result!r}")
        info = result.info
        if state in {"PROGRESS", "STARTED"} and isinstance(info, dict):
            return dict(info)
        return None

    return _wait_until(
        started,
        timeout_seconds=timeout_seconds,
        detail=f"probe {result.id} to start",
    )


def _wait_for_result(result, *, timeout_seconds: float) -> dict[str, Any]:
    value = result.get(timeout=timeout_seconds, propagate=True, disable_sync_subtasks=False)
    if not isinstance(value, dict):
        raise RuntimeError(f"Probe {result.id} returned an invalid result: {value!r}")
    return dict(value)


def _publish(probe_task, probe_id: str, *, mode: str, delay_seconds: float):
    return probe_task.apply_async(
        args=(probe_id,),
        kwargs={
            "mode": mode,
            "delay_seconds": delay_seconds,
            "track_delivery": True,
        },
        queue="default",
    )


def _assert_probe(value: dict[str, Any], *, probe_id: str) -> None:
    if value.get("probe_id") != probe_id:
        raise RuntimeError(f"Probe identity mismatch: {value!r}")


def run_drill(args: argparse.Namespace) -> dict[str, Any]:
    redis_url = str(args.redis_url or "").strip()
    parsed_redis = urlsplit(redis_url)
    if parsed_redis.scheme not in {"redis", "rediss"} or not parsed_redis.hostname:
        raise ValueError("--redis-url must be an explicit redis:// or rediss:// URL")
    if args.visibility_timeout_seconds < 30:
        raise ValueError("--visibility-timeout-seconds must be at least 30")

    run_id = uuid.uuid4().hex
    started_at = datetime.now(UTC)
    with tempfile.TemporaryDirectory(prefix="govoplan-worker-drill-") as directory:
        root = Path(directory)
        database_url = f"sqlite:///{root / 'runtime.db'}"
        environment = dict(os.environ)
        environment.update(
            {
                "REDIS_URL": redis_url,
                "CELERY_ENABLED": "true",
                "CELERY_QUEUES": "default",
                "CELERY_VISIBILITY_TIMEOUT_SECONDS": str(
                    args.visibility_timeout_seconds
                ),
                "DATABASE_URL": database_url,
                "ENABLED_MODULES": "access",
                "APP_ENV": "development",
                "GOVOPLAN_EXPECTED_WORKER_REPLICAS": "1",
                "GOVOPLAN_WORKER_POOL": "acceptance-drill",
            }
        )
        os.environ.update(environment)

        from govoplan_core.db.migrations import migrate_database

        migrate_database(
            database_url=database_url,
            enabled_modules=("access",),
        )
        from govoplan_core.celery_app import celery, worker_acceptance_probe

        celery.backend.client.ping()
        celery.control.purge()
        worker: subprocess.Popen[bytes] | None = None
        worker_log: Any = None
        evidence: dict[str, Any] = {
            "schema_version": "1.0",
            "run_id": run_id,
            "started_at": started_at.isoformat(),
            "redis": _redacted_redis_url(redis_url),
            "visibility_timeout_seconds": args.visibility_timeout_seconds,
            "checks": [],
        }
        try:
            prefix = f"govoplan-drill-{run_id[:8]}-a@"
            worker, worker_log = _start_worker(
                python=args.python,
                hostname=prefix + "%h",
                environment=environment,
                log_path=root / "worker-a.log",
            )
            hostname = _wait_for_worker(
                celery,
                worker,
                prefix,
                args.timeout_seconds,
            )
            evidence["checks"].append(
                {"id": "worker_startup", "state": "passed", "worker": hostname}
            )

            probe_id = f"{run_id}-publish"
            result = _publish(
                worker_acceptance_probe,
                probe_id,
                mode="complete",
                delay_seconds=0,
            )
            value = _wait_for_result(result, timeout_seconds=args.timeout_seconds)
            _assert_probe(value, probe_id=probe_id)
            if value.get("delivery_count") != 1:
                raise RuntimeError(f"Publish probe was not delivered exactly once: {value!r}")
            evidence["checks"].append(
                {
                    "id": "publish_consume",
                    "state": "passed",
                    "task_id": result.id,
                    "delivery_count": value["delivery_count"],
                }
            )

            probe_id = f"{run_id}-retry"
            result = _publish(
                worker_acceptance_probe,
                probe_id,
                mode="retry_once",
                delay_seconds=0,
            )
            value = _wait_for_result(result, timeout_seconds=args.timeout_seconds)
            _assert_probe(value, probe_id=probe_id)
            if value.get("retries") != 1 or value.get("delivery_count") != 2:
                raise RuntimeError(f"Retry probe did not execute twice: {value!r}")
            evidence["checks"].append(
                {
                    "id": "application_retry",
                    "state": "passed",
                    "task_id": result.id,
                    "delivery_count": value["delivery_count"],
                }
            )

            probe_id = f"{run_id}-warm"
            result = _publish(
                worker_acceptance_probe,
                probe_id,
                mode="complete",
                delay_seconds=2,
            )
            _wait_for_started(result, timeout_seconds=args.timeout_seconds)
            worker.send_signal(signal.SIGTERM)
            value = _wait_for_result(result, timeout_seconds=args.timeout_seconds)
            _assert_probe(value, probe_id=probe_id)
            worker.wait(timeout=args.timeout_seconds)
            if worker.returncode != 0 or value.get("delivery_count") != 1:
                raise RuntimeError(
                    f"Warm worker shutdown did not finish in-flight work: {value!r}"
                )
            worker_log.close()
            worker = None
            worker_log = None
            evidence["checks"].append(
                {
                    "id": "graceful_shutdown",
                    "state": "passed",
                    "task_id": result.id,
                    "delivery_count": value["delivery_count"],
                }
            )

            prefix = f"govoplan-drill-{run_id[:8]}-b@"
            worker, worker_log = _start_worker(
                python=args.python,
                hostname=prefix + "%h",
                environment=environment,
                log_path=root / "worker-b.log",
            )
            _wait_for_worker(celery, worker, prefix, args.timeout_seconds)
            probe_id = f"{run_id}-loss"
            result = _publish(
                worker_acceptance_probe,
                probe_id,
                mode="complete",
                delay_seconds=min(120.0, args.visibility_timeout_seconds + 20.0),
            )
            first_started = _wait_for_started(
                result,
                timeout_seconds=args.timeout_seconds,
            )
            if first_started.get("delivery_count") != 1:
                raise RuntimeError(f"Worker-loss probe did not start once: {first_started!r}")
            worker.kill()
            worker.wait(timeout=10)
            worker_log.close()
            worker = None
            worker_log = None

            prefix = f"govoplan-drill-{run_id[:8]}-c@"
            worker, worker_log = _start_worker(
                python=args.python,
                hostname=prefix + "%h",
                environment=environment,
                log_path=root / "worker-c.log",
            )
            _wait_for_worker(celery, worker, prefix, args.timeout_seconds)
            value = _wait_for_result(
                result,
                timeout_seconds=args.timeout_seconds
                + args.visibility_timeout_seconds
                + 30,
            )
            _assert_probe(value, probe_id=probe_id)
            if value.get("delivery_count") != 2:
                raise RuntimeError(f"Worker-loss probe was not redelivered: {value!r}")
            evidence["checks"].append(
                {
                    "id": "worker_loss_redelivery",
                    "state": "passed",
                    "task_id": result.id,
                    "delivery_count": value["delivery_count"],
                    "broker_redelivered": bool(value.get("redelivered")),
                }
            )
        except BaseException as exc:
            evidence["error"] = f"{type(exc).__name__}: {exc}"
            evidence["result"] = {"state": "failed"}
            raise
        finally:
            _stop_worker(worker, worker_log)
            celery.control.purge()
        evidence["completed_at"] = datetime.now(UTC).isoformat()
        evidence["result"] = {"state": "passed"}
        return evidence


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        evidence = run_drill(args)
    except (OSError, RuntimeError, TimeoutError, ValueError) as exc:
        print(f"worker runtime drill failed: {exc}", file=sys.stderr)
        return 1
    payload = json.dumps(evidence, indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        temporary = args.output.with_suffix(args.output.suffix + ".tmp")
        temporary.write_text(payload, encoding="utf-8")
        temporary.replace(args.output)
        print(f"Worker runtime evidence written to {args.output}")
    else:
        print(payload, end="")
    return 0


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