from __future__ import annotations

import json
import os
import queue
import signal
import subprocess
import threading
import time
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from pathlib import Path

import shutil
import uuid

import codex_state
from codex_appserver import CodexAppServer
from health_store import HealthSnapshotStore


class HealthStatus(Enum):
    OK = "ok"
    BROKEN = "broken"
    UNKNOWN = "unknown"


@dataclass(frozen=True)
class NamedLimit:
    kind: str
    group: str | None
    percent: int | None
    resets_at: float | None
    active: bool


@dataclass(frozen=True)
class ExtraUsageSummary:
    used: float | None = None
    limit: float | None = None
    unit: str | None = None
    display: str | None = None


@dataclass(frozen=True)
class SpendSummary:
    amount: float | None = None
    limit: float | None = None
    currency: str | None = None
    display: str | None = None


@dataclass(frozen=True)
class AccountSnapshot:
    status: HealthStatus
    primary_used_pct: int | None
    secondary_used_pct: int | None
    primary_reset_at: float | None = None
    secondary_reset_at: float | None = None
    checked_at: float | None = None
    detail: str | None = None
    named_limits: tuple[NamedLimit, ...] = ()
    extra_usage: ExtraUsageSummary | None = None
    spend: SpendSummary | None = None


def codex_executable() -> str | None:
    for directory in os.environ.get("PATH", "").split(os.pathsep):
        candidate = Path(directory or ".") / "codex"
        if not candidate.is_file() or not os.access(candidate, os.X_OK):
            continue
        try:
            resolved = candidate.resolve()
        except OSError:
            continue
        if resolved.name == "_tmpjail-shim.sh":
            continue
        return str(resolved)
    return None


class AccountHealthClient:
    _scope_available: bool | None = None

    def __init__(self) -> None:
        self._protocol = CodexAppServer()

    @staticmethod
    def _codex_executable() -> str | None:
        return codex_executable()

    @classmethod
    def _user_scope_available(cls) -> bool:
        if cls._scope_available is None:
            if shutil.which("systemd-run") is None or shutil.which("systemctl") is None:
                cls._scope_available = False
            else:
                try:
                    probe = subprocess.run(
                        ["systemctl", "--user", "show-environment"],
                        stdout=subprocess.DEVNULL,
                        stderr=subprocess.DEVNULL,
                        timeout=5,
                    )
                    cls._scope_available = probe.returncode == 0
                except (OSError, subprocess.TimeoutExpired):
                    cls._scope_available = False
        return cls._scope_available

    def fetch(
        self,
        codex_home: Path,
        timeout_secs: float = 10.0,
        *,
        allow_warmup: bool = False,
    ) -> AccountSnapshot:
        if not codex_state.needs_warmup(codex_home):
            return self._probe(codex_home, timeout_secs)
        if not allow_warmup or codex_state.warmup_in_progress(codex_home):
            return AccountSnapshot(
                HealthStatus.UNKNOWN, None, None, detail="preparing account data"
            )
        codex_state.release_orphaned_claim(codex_home)
        with codex_state.warmup_claim(codex_home):
            return self._probe(codex_home, max(timeout_secs, codex_state.WARMUP_TIMEOUT_S))

    def _probe(self, codex_home: Path, timeout_secs: float) -> AccountSnapshot:
        executable = self._codex_executable()
        if executable is None:
            return AccountSnapshot(
                HealthStatus.UNKNOWN, None, None, detail="Codex runtime unavailable"
            )
        argv = [executable, "app-server", "--stdio"]
        scope_unit: str | None = None
        if self._user_scope_available():
            # `codex` resolves through wrapper layers (shim -> tmpjail -> node
            # launcher -> vendor ELF) that can detach from the process group
            # (39 leaked app-servers, 2026-08-04). An owned scope makes teardown
            # a cgroup-wide kill instead of a pgroup guess.
            scope_unit = f"codex-probe-{os.getpid()}-{uuid.uuid4().hex[:8]}"
            argv = [
                "systemd-run",
                "--user",
                "--scope",
                "--collect",
                "--quiet",
                "--slice=agent.slice",
                f"--unit={scope_unit}",
                # A scope with no ceiling inherits agent.slice's, so one runaway probe
                # can consume the whole slice. Measured peak is 58 MiB.
                "-p",
                "MemoryMax=2G",
                "-p",
                "MemoryHigh=1G",
                "-p",
                "TasksMax=512",
                "--",
                *argv,
            ]
        try:
            process = subprocess.Popen(
                argv,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                env={**os.environ, "CODEX_HOME": str(codex_home)},
                start_new_session=True,
            )
        except OSError:
            return AccountSnapshot(HealthStatus.UNKNOWN, None, None)

        try:
            if process.stdin is None or process.stdout is None:
                return AccountSnapshot(HealthStatus.UNKNOWN, None, None)

            initialize_line, read_line = self._protocol.request_lines()
            process.stdin.write(initialize_line)
            process.stdin.write(read_line)
            process.stdin.flush()

            line_queue: queue.Queue[str | None] = queue.Queue()

            def _read_lines() -> None:
                try:
                    while True:
                        chunk = process.stdout.readline()
                        line_queue.put(chunk)
                        if not chunk:
                            return
                except (OSError, ValueError):
                    line_queue.put(None)

            reader_thread = threading.Thread(target=_read_lines, daemon=True)
            reader_thread.start()

            deadline = time.monotonic() + timeout_secs
            while True:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    return AccountSnapshot(HealthStatus.UNKNOWN, None, None)

                try:
                    line = line_queue.get(timeout=remaining)
                except queue.Empty:
                    return AccountSnapshot(HealthStatus.UNKNOWN, None, None)

                if line is None:
                    return AccountSnapshot(HealthStatus.UNKNOWN, None, None)

                if not line:
                    if process.poll() is not None:
                        return AccountSnapshot(HealthStatus.UNKNOWN, None, None)
                    continue

                reply = self._protocol.parse_rate_limits_response(line)
                if reply is None:
                    continue

                if reply.is_error:
                    return AccountSnapshot(HealthStatus.BROKEN, None, None)

                payload = json.loads(line)
                result = payload.get("result")
                if not isinstance(result, dict):
                    return AccountSnapshot(HealthStatus.UNKNOWN, None, None)

                rate_limits = self._rate_limits_payload(result)
                if not ({"primary", "secondary"} & rate_limits.keys()):
                    return AccountSnapshot(HealthStatus.UNKNOWN, None, None)
                now = time.time()
                primary_window, secondary_window = self._semantic_windows(rate_limits)
                primary = self._used_percent(primary_window)
                secondary = self._used_percent(secondary_window)
                return AccountSnapshot(
                    HealthStatus.OK,
                    primary,
                    secondary,
                    primary_reset_at=self._reset_at(primary_window, now),
                    secondary_reset_at=self._reset_at(secondary_window, now),
                )
        finally:
            self._terminate_process(process, scope_unit)

    def write_cache(self, cache_path: Path, snapshots: dict[str, AccountSnapshot]) -> None:
        HealthSnapshotStore(cache_path, stale_after_s=0).write(snapshots)

    def read_cache(self, cache_path: Path) -> dict[str, AccountSnapshot]:
        return HealthSnapshotStore(cache_path, stale_after_s=0).read()

    @staticmethod
    def _used_percent(window: object) -> int | None:
        if not isinstance(window, dict):
            return None
        value = window.get("usedPercent")
        return value if isinstance(value, int) else None

    @staticmethod
    def _rate_limits_payload(result: object) -> dict[str, object]:
        if not isinstance(result, dict):
            return {}
        nested = result.get("rateLimits")
        if isinstance(nested, dict):
            return nested
        return result

    @staticmethod
    def _semantic_windows(rate_limits: dict[str, object]) -> tuple[object, object]:
        positional = (rate_limits.get("primary"), rate_limits.get("secondary"))
        durations = tuple(AccountHealthClient._window_duration_mins(window) for window in positional)
        if all(duration is None for duration in durations):
            return positional

        primary_window = None
        secondary_window = None
        for window, duration in zip(positional, durations, strict=True):
            if duration == 300:
                primary_window = window
            elif duration == 10_080:
                secondary_window = window
        return primary_window, secondary_window

    @staticmethod
    def _window_duration_mins(window: object) -> int | None:
        if not isinstance(window, dict):
            return None
        value = window.get("windowDurationMins")
        if isinstance(value, bool) or not isinstance(value, (int, float)):
            return None
        return int(value)

    @staticmethod
    def _reset_at(window: object, now: float) -> float | None:
        if not isinstance(window, dict):
            return None

        for key in (
            "resetSeconds",
            "resetsInSeconds",
            "secondsUntilReset",
            "resetInSeconds",
        ):
            value = AccountHealthClient._numeric_value(window.get(key))
            if value is not None:
                return now + max(0.0, value)

        for key in ("resetsAt", "resetAt", "resetTime", "resetTimestamp"):
            value = window.get(key)
            numeric = AccountHealthClient._numeric_value(value)
            if numeric is not None:
                if numeric > 10_000_000_000:
                    return numeric / 1000.0
                if numeric < 1_000_000_000:
                    return now + max(0.0, numeric)
                return numeric
            if isinstance(value, str):
                parsed = AccountHealthClient._parse_iso_timestamp(value)
                if parsed is not None:
                    return parsed
        return None

    @staticmethod
    def _numeric_value(value: object) -> float | None:
        if isinstance(value, bool):
            return None
        if isinstance(value, (int, float)):
            return float(value)
        if isinstance(value, str):
            try:
                return float(value)
            except ValueError:
                return None
        return None

    @staticmethod
    def _parse_iso_timestamp(value: str) -> float | None:
        normalized = value.strip()
        if not normalized:
            return None
        if normalized.endswith("Z"):
            normalized = f"{normalized[:-1]}+00:00"
        try:
            return datetime.fromisoformat(normalized).timestamp()
        except ValueError:
            return None

    @staticmethod
    def _kill_scope(scope_unit: str) -> None:
        try:
            subprocess.run(
                [
                    "systemctl",
                    "--user",
                    "kill",
                    "--kill-whom=all",
                    "--signal=SIGKILL",
                    f"{scope_unit}.scope",
                ],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                timeout=5,
            )
        except (OSError, subprocess.TimeoutExpired):
            return

    @staticmethod
    def _terminate_process(
        process: subprocess.Popen[str], scope_unit: str | None = None
    ) -> None:
        def _signal_group(sig: int) -> None:
            try:
                os.killpg(process.pid, sig)
            except (OSError, ProcessLookupError):
                process.send_signal(sig)

        try:
            try:
                _signal_group(signal.SIGTERM)
            except OSError:
                return

            try:
                process.wait(timeout=1.0)
            except (OSError, subprocess.TimeoutExpired):
                try:
                    _signal_group(signal.SIGKILL)
                except OSError:
                    return
                try:
                    process.wait(timeout=1.0)
                except (OSError, subprocess.TimeoutExpired):
                    return
        finally:
            if scope_unit is not None:
                AccountHealthClient._kill_scope(scope_unit)