"""Node lifecycle operations using the Kubernetes API.

Provides cordon, uncordon, and drain via the Eviction API (not kubectl CLI).
The Eviction API respects PodDisruptionBudgets, unlike `kubectl drain --force`.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING

from kubernetes import client, config
from kubernetes.client.exceptions import ApiException

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)

# Node label schema for Citadel lifecycle management
LABEL_TERMINATE_REASON = "citadel.dev/terminate-reason"
LABEL_TERMINATE_AFTER = "citadel.dev/terminate-after"
LABEL_POLICY_NAME = "citadel.dev/policy-name"
LABEL_MANAGED_BY = "citadel.dev/managed-by"

MANAGED_BY_VALUE = "citadel-lifecycle"


@dataclass(frozen=True)
class NodeInfo:
    """Immutable snapshot of a node's lifecycle-relevant state."""

    name: str
    creation_timestamp: datetime
    age_hours: float
    ready: bool
    schedulable: bool
    conditions: dict[str, str] = field(default_factory=dict)
    labels: dict[str, str] = field(default_factory=dict)

    @property
    def marked_for_termination(self) -> bool:
        return LABEL_TERMINATE_REASON in self.labels

    @property
    def terminate_reason(self) -> str | None:
        return self.labels.get(LABEL_TERMINATE_REASON)


@dataclass(frozen=True)
class DrainResult:
    """Result of a node drain operation."""

    node_name: str
    success: bool
    pods_evicted: int
    pods_failed: list[str] = field(default_factory=list)
    pods_force_deleted: list[str] = field(default_factory=list)
    error: str | None = None


class NodeManager:
    """Manages Kubernetes node lifecycle operations.

    Uses the Eviction API for draining, which respects PodDisruptionBudgets.
    This is the substrate that Apollo's spoke controller calls into.
    """

    def __init__(self, kubeconfig: str | None = None) -> None:
        if kubeconfig:
            config.load_kube_config(config_file=kubeconfig)
        else:
            try:
                config.load_incluster_config()
            except config.ConfigException:
                config.load_kube_config()

        self._core = client.CoreV1Api()
        self._policy = client.PolicyV1Api()

    def list_nodes(self) -> list[NodeInfo]:
        """List all nodes with lifecycle-relevant info."""
        nodes = self._core.list_node()
        return [self._to_node_info(n) for n in nodes.items]

    def get_node(self, name: str) -> NodeInfo:
        """Get a single node's lifecycle info."""
        node = self._core.read_node(name)
        return self._to_node_info(node)

    def get_node_age_hours(self, name: str) -> float:
        """Get the age of a node in hours."""
        info = self.get_node(name)
        return info.age_hours

    def cordon(self, name: str) -> None:
        """Mark a node as unschedulable (cordon).

        Prevents new pods from being scheduled on this node.
        Existing pods continue running.
        """
        body = {"spec": {"unschedulable": True}}
        self._core.patch_node(name, body)
        logger.info("Cordoned node %s", name)

    def uncordon(self, name: str) -> None:
        """Mark a node as schedulable (uncordon)."""
        body = {"spec": {"unschedulable": False}}
        self._core.patch_node(name, body)
        logger.info("Uncordoned node %s", name)

    def label_for_termination(
        self, name: str, reason: str, policy_name: str
    ) -> None:
        """Label a node with termination intent.

        This records *why* the node is being terminated and *which policy*
        triggered it, creating an auditable trail before any destructive
        action occurs. Apollo's spoke checks these labels.
        """
        body = {
            "metadata": {
                "labels": {
                    LABEL_TERMINATE_REASON: reason,
                    LABEL_POLICY_NAME: policy_name,
                    LABEL_MANAGED_BY: MANAGED_BY_VALUE,
                    LABEL_TERMINATE_AFTER: datetime.now(timezone.utc).strftime(
                        "%Y-%m-%dT%H-%M-%SZ"
                    ),
                }
            }
        }
        self._core.patch_node(name, body)
        logger.info(
            "Labeled node %s for termination: reason=%s policy=%s",
            name,
            reason,
            policy_name,
        )

    def clear_termination_labels(self, name: str) -> None:
        """Remove termination labels from a node."""
        body = {
            "metadata": {
                "labels": {
                    LABEL_TERMINATE_REASON: None,
                    LABEL_POLICY_NAME: None,
                    LABEL_MANAGED_BY: None,
                    LABEL_TERMINATE_AFTER: None,
                }
            }
        }
        self._core.patch_node(name, body)
        logger.info("Cleared termination labels from node %s", name)

    def drain(
        self,
        name: str,
        timeout_seconds: int = 300,
        delete_emptydir_data: bool = True,
        force_after_timeout: int | None = None,
    ) -> DrainResult:
        """Drain a node using the Eviction API.

        Unlike `kubectl drain --force`, this uses the Eviction API which
        respects PodDisruptionBudgets. If a PDB prevents eviction, the
        pod is retried until timeout rather than force-deleted.

        When ``force_after_timeout`` is set and the graceful eviction timeout
        expires, remaining pods are force-deleted directly via the Kubernetes
        API, bypassing PDB constraints.  This mirrors Palantir's behaviour
        where a stuck PDB must not block node replacement indefinitely.

        Args:
            name: Node name to drain.
            timeout_seconds: Max time to wait for all pods to evict.
            delete_emptydir_data: Whether to evict pods using emptyDir volumes.
            force_after_timeout: Seconds after which to force-delete pods that
                could not be evicted.  ``None`` (default) disables force drain.

        Returns:
            DrainResult with eviction details.
        """
        # Step 1: Cordon first
        self.cordon(name)

        # Step 2: Get all non-daemonset pods on this node
        field_selector = f"spec.nodeName={name}"
        pods = self._core.list_pod_for_all_namespaces(
            field_selector=field_selector
        )

        pods_to_evict = []
        for pod in pods.items:
            # Skip mirror pods (static pods managed by kubelet)
            if pod.metadata.annotations and pod.metadata.annotations.get(
                "kubernetes.io/config.mirror"
            ):
                continue

            # Skip DaemonSet-managed pods
            if self._is_daemonset_pod(pod):
                continue

            # Skip pods using emptyDir if not deleting that data
            if not delete_emptydir_data and self._uses_emptydir(pod):
                continue

            pods_to_evict.append(pod)

        logger.info(
            "Draining node %s: %d pods to evict", name, len(pods_to_evict)
        )

        # Step 3: Evict pods via Eviction API
        pods_evicted = 0
        pods_failed: list[str] = []
        deadline = time.monotonic() + timeout_seconds

        for pod in pods_to_evict:
            pod_name = pod.metadata.name
            pod_ns = pod.metadata.namespace
            pod_ref = f"{pod_ns}/{pod_name}"

            evicted = self._evict_pod_with_retry(
                pod_name, pod_ns, deadline
            )
            if evicted:
                pods_evicted += 1
            else:
                pods_failed.append(pod_ref)
                logger.warning("Failed to evict pod %s", pod_ref)

        # Step 4: Force-delete remaining pods if configured
        pods_force_deleted: list[str] = []
        if pods_failed and force_after_timeout is not None:
            logger.warning(
                "Force-draining node %s: %d pod(s) could not be evicted within "
                "timeout, force-deleting after %ds (bypasses PDB constraints)",
                name,
                len(pods_failed),
                force_after_timeout,
            )
            force_deadline = time.monotonic() + force_after_timeout
            still_failed: list[str] = []

            for pod_ref in pods_failed:
                pod_ns, pod_name = pod_ref.split("/", 1)
                deleted = self._force_delete_pod(
                    pod_name, pod_ns, force_deadline
                )
                if deleted:
                    pods_force_deleted.append(pod_ref)
                else:
                    still_failed.append(pod_ref)

            pods_failed = still_failed

        success = len(pods_failed) == 0
        error_msg = None
        if pods_failed:
            error_msg = (
                f"Failed to evict {len(pods_failed)} pod(s): "
                f"{', '.join(pods_failed)}"
            )

        result = DrainResult(
            node_name=name,
            success=success,
            pods_evicted=pods_evicted,
            pods_failed=pods_failed,
            pods_force_deleted=pods_force_deleted,
            error=error_msg,
        )

        if success and not pods_force_deleted:
            logger.info(
                "Successfully drained node %s (%d pods evicted)",
                name,
                pods_evicted,
            )
        elif success and pods_force_deleted:
            logger.warning(
                "Drained node %s (%d evicted, %d force-deleted)",
                name,
                pods_evicted,
                len(pods_force_deleted),
            )
        else:
            logger.warning(
                "Drain of node %s completed with errors: %s",
                name,
                error_msg,
            )

        return result

    def _force_delete_pod(
        self, name: str, namespace: str, deadline: float
    ) -> bool:
        """Force-delete a pod directly, bypassing PDB constraints.

        Uses grace_period_seconds=0 to request immediate termination.
        This is the nuclear option when a PDB blocks eviction past an
        acceptable threshold.
        """
        try:
            self._core.delete_namespaced_pod(
                name=name,
                namespace=namespace,
                body=client.V1DeleteOptions(grace_period_seconds=0),
            )
            logger.warning(
                "Force-deleted pod %s/%s (PDB bypassed)", namespace, name
            )
            self._wait_pod_deleted(name, namespace, deadline)
            return True
        except ApiException as e:
            if e.status == 404:
                # Pod already gone
                return True
            logger.error(
                "Failed to force-delete pod %s/%s: %s",
                namespace,
                name,
                e.reason,
            )
            return False

    def _evict_pod_with_retry(
        self, name: str, namespace: str, deadline: float
    ) -> bool:
        """Evict a single pod, retrying if PDB blocks it."""
        eviction = client.V1Eviction(
            metadata=client.V1ObjectMeta(name=name, namespace=namespace),
            delete_options=client.V1DeleteOptions(
                grace_period_seconds=30
            ),
        )

        while time.monotonic() < deadline:
            try:
                self._core.create_namespaced_pod_eviction(
                    name=name, namespace=namespace, body=eviction
                )
                # Wait for pod to actually terminate
                self._wait_pod_deleted(name, namespace, deadline)
                return True
            except ApiException as e:
                if e.status == 429:
                    # 429 Too Many Requests = PDB is blocking eviction
                    logger.debug(
                        "PDB blocking eviction of %s/%s, retrying...",
                        namespace,
                        name,
                    )
                    time.sleep(5)
                    continue
                elif e.status == 404:
                    # Pod already gone
                    return True
                else:
                    logger.error(
                        "Eviction API error for %s/%s: %s",
                        namespace,
                        name,
                        e.reason,
                    )
                    return False

        logger.warning(
            "Eviction of %s/%s timed out (PDB may be blocking)",
            namespace,
            name,
        )
        return False

    def _wait_pod_deleted(
        self, name: str, namespace: str, deadline: float
    ) -> None:
        """Wait for a pod to be fully deleted."""
        while time.monotonic() < deadline:
            try:
                pod = self._core.read_namespaced_pod(name, namespace)
                if pod.metadata.deletion_timestamp is not None:
                    time.sleep(2)
                    continue
                time.sleep(2)
            except ApiException as e:
                if e.status == 404:
                    return
                raise

    def _to_node_info(self, node: client.V1Node) -> NodeInfo:
        """Convert a Kubernetes Node object to NodeInfo."""
        creation = node.metadata.creation_timestamp
        now = datetime.now(timezone.utc)
        age_hours = (now - creation).total_seconds() / 3600.0

        conditions = {}
        ready = False
        for cond in node.status.conditions or []:
            conditions[cond.type] = cond.status
            if cond.type == "Ready":
                ready = cond.status == "True"

        schedulable = not (node.spec.unschedulable or False)

        labels = dict(node.metadata.labels or {})

        return NodeInfo(
            name=node.metadata.name,
            creation_timestamp=creation,
            age_hours=round(age_hours, 2),
            ready=ready,
            schedulable=schedulable,
            conditions=conditions,
            labels=labels,
        )

    @staticmethod
    def _is_daemonset_pod(pod: client.V1Pod) -> bool:
        """Check if a pod is managed by a DaemonSet."""
        for ref in pod.metadata.owner_references or []:
            if ref.kind == "DaemonSet":
                return True
        return False

    @staticmethod
    def _uses_emptydir(pod: client.V1Pod) -> bool:
        """Check if a pod uses emptyDir volumes."""
        for vol in pod.spec.volumes or []:
            if vol.empty_dir is not None:
                return True
        return False
