"""Custom probe checker for detecting node-level failure modes.

Extends the standard Kubernetes condition-based health checking with
custom exec and HTTP probes. Inspired by Palantir's custom bad-state
detectors that detect stuck volumes, GPU errors, NVMe degradation, etc.

Each probe tracks consecutive failures independently. A probe is only
considered "failed" once it exceeds its configured failureThreshold,
preventing transient blips from triggering node termination.
"""

from __future__ import annotations

import logging
import subprocess
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from enum import Enum

logger = logging.getLogger(__name__)


class ProbeType(Enum):
    """Supported probe types."""

    EXEC = "exec"
    HTTP = "http"


class ProbeStatus(Enum):
    """Current status of a probe after evaluation."""

    PASSING = "passing"
    FAILING = "failing"
    UNKNOWN = "unknown"


@dataclass(frozen=True)
class ExecProbeConfig:
    """Configuration for an exec-type probe."""

    command: list[str]
    timeout_seconds: int = 10


@dataclass(frozen=True)
class HttpProbeConfig:
    """Configuration for an HTTP-type probe."""

    path: str
    port: int
    scheme: str = "http"
    host: str = "localhost"

    @property
    def url(self) -> str:
        """Build the full probe URL."""
        return f"{self.scheme}://{self.host}:{self.port}{self.path}"


@dataclass(frozen=True)
class CustomProbeSpec:
    """Full specification for a custom probe, parsed from the CRD."""

    name: str
    probe_type: ProbeType
    failure_threshold: int = 3
    period_seconds: int = 60
    exec_config: ExecProbeConfig | None = None
    http_config: HttpProbeConfig | None = None

    def __post_init__(self) -> None:
        if self.probe_type == ProbeType.EXEC and self.exec_config is None:
            raise ValueError(
                f"Probe '{self.name}': exec probe requires exec_config"
            )
        if self.probe_type == ProbeType.HTTP and self.http_config is None:
            raise ValueError(
                f"Probe '{self.name}': http probe requires http_config"
            )


@dataclass(frozen=True)
class ProbeResult:
    """Result of a single probe execution."""

    probe_name: str
    probe_type: ProbeType
    status: ProbeStatus
    consecutive_failures: int
    failure_threshold: int
    message: str
    timestamp: float

    @property
    def threshold_exceeded(self) -> bool:
        """True when consecutive failures have reached the failure threshold."""
        return self.consecutive_failures >= self.failure_threshold


@dataclass(frozen=True)
class CustomProbeReport:
    """Aggregated results for all custom probes on a node."""

    node_name: str
    results: list[ProbeResult]
    timestamp: float

    @property
    def any_failed(self) -> bool:
        """True if any probe has exceeded its failure threshold."""
        return any(r.threshold_exceeded for r in self.results)

    @property
    def failed_probes(self) -> list[ProbeResult]:
        """Return only the probes that have exceeded their failure threshold."""
        return [r for r in self.results if r.threshold_exceeded]

    @property
    def passing_probes(self) -> list[ProbeResult]:
        """Return only the probes that are currently passing."""
        return [r for r in self.results if r.status == ProbeStatus.PASSING]


def parse_custom_probes(raw_probes: list[dict]) -> list[CustomProbeSpec]:
    """Parse custom probe specs from the CRD's raw YAML/JSON representation.

    Args:
        raw_probes: List of probe dicts as they appear in the CRD spec.

    Returns:
        List of validated CustomProbeSpec objects.

    Raises:
        ValueError: If a probe spec is invalid.
    """
    specs: list[CustomProbeSpec] = []

    for raw in raw_probes:
        name = raw["name"]
        probe_type = ProbeType(raw["type"])
        failure_threshold = raw.get("failureThreshold", 3)
        period_seconds = raw.get("periodSeconds", 60)

        exec_config = None
        http_config = None

        if probe_type == ProbeType.EXEC:
            exec_raw = raw.get("exec", {})
            exec_config = ExecProbeConfig(
                command=exec_raw["command"],
                timeout_seconds=exec_raw.get("timeoutSeconds", 10),
            )
        elif probe_type == ProbeType.HTTP:
            http_raw = raw.get("http", {})
            http_config = HttpProbeConfig(
                path=http_raw["path"],
                port=http_raw["port"],
                scheme=http_raw.get("scheme", "http"),
                host=http_raw.get("host", "localhost"),
            )

        specs.append(
            CustomProbeSpec(
                name=name,
                probe_type=probe_type,
                failure_threshold=failure_threshold,
                period_seconds=period_seconds,
                exec_config=exec_config,
                http_config=http_config,
            )
        )

    return specs


class CustomProbeChecker:
    """Executes custom probes and tracks consecutive failures per probe per node.

    This is the detection layer for failure modes that Kubernetes conditions
    do not cover: stuck volumes, GPU health, NVMe degradation, NUMA
    imbalances, etc.

    The checker maintains state across invocations to track consecutive
    failures. A probe only reports as failed once it exceeds its configured
    failureThreshold, filtering out transient errors.

    Usage::

        checker = CustomProbeChecker()
        probes = parse_custom_probes(policy_spec["selector"]["customProbes"])
        report = checker.check(node_name="worker-1", probes=probes)
        if report.any_failed:
            for p in report.failed_probes:
                logger.warning("Probe %s failed on %s: %s", p.probe_name, node_name, p.message)
    """

    def __init__(self) -> None:
        # Tracks consecutive failures: {(node_name, probe_name): count}
        self._failure_counts: dict[tuple[str, str], int] = {}
        # Tracks last execution time: {(node_name, probe_name): timestamp}
        self._last_run: dict[tuple[str, str], float] = {}

    def check(
        self,
        node_name: str,
        probes: list[CustomProbeSpec],
    ) -> CustomProbeReport:
        """Execute all probes for a node and return an aggregated report.

        Probes that have not reached their periodSeconds interval since
        the last run are skipped (their previous result is carried forward).

        Args:
            node_name: Name of the node being checked.
            probes: List of probe specifications to execute.

        Returns:
            CustomProbeReport with results for all probes.
        """
        now = time.monotonic()
        results: list[ProbeResult] = []

        for probe in probes:
            key = (node_name, probe.name)
            last = self._last_run.get(key, 0.0)

            # Skip probes not yet due
            if now - last < probe.period_seconds and key in self._failure_counts:
                # Carry forward the previous state
                count = self._failure_counts[key]
                status = ProbeStatus.FAILING if count > 0 else ProbeStatus.PASSING
                results.append(
                    ProbeResult(
                        probe_name=probe.name,
                        probe_type=probe.probe_type,
                        status=status,
                        consecutive_failures=count,
                        failure_threshold=probe.failure_threshold,
                        message="Skipped (period not elapsed)",
                        timestamp=last,
                    )
                )
                continue

            # Execute the probe
            success, message = self._execute_probe(probe)
            self._last_run[key] = now

            if success:
                self._failure_counts[key] = 0
                results.append(
                    ProbeResult(
                        probe_name=probe.name,
                        probe_type=probe.probe_type,
                        status=ProbeStatus.PASSING,
                        consecutive_failures=0,
                        failure_threshold=probe.failure_threshold,
                        message=message,
                        timestamp=now,
                    )
                )
            else:
                count = self._failure_counts.get(key, 0) + 1
                self._failure_counts[key] = count
                status = (
                    ProbeStatus.FAILING
                    if count >= probe.failure_threshold
                    else ProbeStatus.UNKNOWN
                )
                results.append(
                    ProbeResult(
                        probe_name=probe.name,
                        probe_type=probe.probe_type,
                        status=status,
                        consecutive_failures=count,
                        failure_threshold=probe.failure_threshold,
                        message=message,
                        timestamp=now,
                    )
                )

        return CustomProbeReport(
            node_name=node_name,
            results=results,
            timestamp=now,
        )

    def reset(self, node_name: str | None = None) -> None:
        """Reset failure tracking state.

        Args:
            node_name: If provided, only reset state for this node.
                       If None, reset all state.
        """
        if node_name is None:
            self._failure_counts.clear()
            self._last_run.clear()
        else:
            keys_to_remove = [
                k for k in self._failure_counts if k[0] == node_name
            ]
            for k in keys_to_remove:
                del self._failure_counts[k]
                self._last_run.pop(k, None)

    def get_failure_count(self, node_name: str, probe_name: str) -> int:
        """Get the current consecutive failure count for a probe on a node."""
        return self._failure_counts.get((node_name, probe_name), 0)

    def _execute_probe(self, probe: CustomProbeSpec) -> tuple[bool, str]:
        """Execute a single probe and return (success, message).

        Dispatches to the appropriate handler based on probe type.
        """
        if probe.probe_type == ProbeType.EXEC:
            return self._execute_exec_probe(probe)
        elif probe.probe_type == ProbeType.HTTP:
            return self._execute_http_probe(probe)
        else:
            return False, f"Unknown probe type: {probe.probe_type}"

    def _execute_exec_probe(self, probe: CustomProbeSpec) -> tuple[bool, str]:
        """Execute an exec probe via subprocess.

        A zero exit code indicates success; anything else is a failure.
        The command's stderr is captured in the failure message to aid
        debugging.
        """
        assert probe.exec_config is not None
        config = probe.exec_config

        try:
            result = subprocess.run(
                config.command,
                capture_output=True,
                text=True,
                timeout=config.timeout_seconds,
                check=False,
            )
            if result.returncode == 0:
                return True, f"Probe '{probe.name}' passed"
            else:
                stderr = result.stderr.strip()[:200] if result.stderr else ""
                return (
                    False,
                    f"Probe '{probe.name}' failed (exit={result.returncode})"
                    + (f": {stderr}" if stderr else ""),
                )
        except subprocess.TimeoutExpired:
            return (
                False,
                f"Probe '{probe.name}' timed out after {config.timeout_seconds}s",
            )
        except FileNotFoundError:
            return (
                False,
                f"Probe '{probe.name}' command not found: {config.command[0]}",
            )
        except OSError as e:
            return False, f"Probe '{probe.name}' OS error: {e}"

    def _execute_http_probe(self, probe: CustomProbeSpec) -> tuple[bool, str]:
        """Execute an HTTP probe via urllib.

        A 2xx response indicates success; anything else is a failure.
        Uses a short timeout to avoid blocking the evaluation loop.
        """
        assert probe.http_config is not None
        config = probe.http_config
        url = config.url

        try:
            request = urllib.request.Request(url, method="GET")
            with urllib.request.urlopen(request, timeout=10) as response:
                status_code = response.status
                if 200 <= status_code < 300:
                    return True, f"Probe '{probe.name}' passed (HTTP {status_code})"
                else:
                    return (
                        False,
                        f"Probe '{probe.name}' failed (HTTP {status_code})",
                    )
        except urllib.error.HTTPError as e:
            return (
                False,
                f"Probe '{probe.name}' failed (HTTP {e.code})",
            )
        except urllib.error.URLError as e:
            return (
                False,
                f"Probe '{probe.name}' connection failed: {e.reason}",
            )
        except TimeoutError:
            return (
                False,
                f"Probe '{probe.name}' timed out connecting to {url}",
            )
        except OSError as e:
            return False, f"Probe '{probe.name}' error: {e}"
