"""Pydantic models for CRIU checkpoint/restore management."""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Annotated

from pydantic import BaseModel, ConfigDict, Field


class CheckpointStatus(str, Enum):
    """Status of a checkpoint/restore operation."""

    PENDING = "pending"
    FREEZING = "freezing"
    CHECKPOINTING = "checkpointing"
    TRANSFERRING = "transferring"  # For migration
    RESTORING = "restoring"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class CheckpointType(str, Enum):
    """Type of checkpoint to create."""

    FULL = "full"  # Full CRIU checkpoint with all state
    MEMORY_ONLY = "memory_only"  # Memory and process state only
    ITERATIVE = "iterative"  # Iterative/incremental checkpoint


class RuntimeType(str, Enum):
    """Runtime types that support checkpointing."""

    DOCKER = "docker"
    NSPAWN = "nspawn"


class CheckpointMetadata(BaseModel):
    """Metadata for a completed checkpoint."""

    model_config = ConfigDict(frozen=True)

    checkpoint_id: str = Field(..., description="Unique identifier (UUID)")
    container_name: str
    checkpoint_name: str = Field(..., description="Human-readable name")
    checkpoint_type: CheckpointType = CheckpointType.FULL

    created_at: datetime
    completed_at: datetime | None = None

    size_bytes: int = 0
    image_dir: str = Field(..., description="Path to CRIU image directory")

    # Runtime information
    runtime: RuntimeType
    container_id: str | None = Field(
        default=None, description="Runtime-specific container ID"
    )

    # Container state at checkpoint time
    container_was_running: bool = True
    pid_at_checkpoint: int | None = Field(
        default=None, description="Leader PID when checkpoint was taken"
    )

    # Verification
    checksum: str | None = Field(default=None, description="SHA256 hash of checkpoint")
    verified: bool = False

    # For iterative checkpoints
    is_iterative: bool = False
    parent_checkpoint_id: str | None = Field(
        default=None, description="Previous checkpoint this is based on"
    )

    # Additional metadata
    leave_running: bool = Field(
        default=False, description="Whether container kept running after checkpoint"
    )
    tcp_established: bool = Field(
        default=False, description="Whether TCP connections were preserved"
    )
    external_unix_sockets: bool = Field(
        default=False, description="Whether external unix sockets were allowed"
    )


class CheckpointResult(BaseModel):
    """Result of a checkpoint operation."""

    model_config = ConfigDict(frozen=False)

    container_name: str
    checkpoint_name: str
    checkpoint_id: str
    started_at: datetime = Field(default_factory=datetime.now)
    completed_at: datetime | None = None
    status: CheckpointStatus = CheckpointStatus.PENDING
    checkpoint_type: CheckpointType = CheckpointType.FULL

    size_bytes: int = 0
    image_dir: str | None = None
    checksum: str | None = None

    # Container state
    container_was_running: bool = True
    pid_at_checkpoint: int | None = None
    leave_running: bool = False

    error_message: str | None = None
    error_output: str | None = None

    @property
    def success(self) -> bool:
        """Return True if checkpoint completed successfully."""
        return self.status == CheckpointStatus.COMPLETED

    @property
    def duration_seconds(self) -> float | None:
        """Calculate duration of checkpoint operation."""
        if self.completed_at is None:
            return None
        return (self.completed_at - self.started_at).total_seconds()


class RestoreResult(BaseModel):
    """Result of a restore operation."""

    model_config = ConfigDict(frozen=False)

    container_name: str
    checkpoint_name: str
    checkpoint_id: str
    started_at: datetime = Field(default_factory=datetime.now)
    completed_at: datetime | None = None
    status: CheckpointStatus = CheckpointStatus.PENDING

    # New container state after restore
    new_container_id: str | None = Field(
        default=None, description="Container ID after restore (may differ from original)"
    )
    new_pid: int | None = Field(
        default=None, description="New leader PID after restore"
    )

    error_message: str | None = None
    error_output: str | None = None

    @property
    def success(self) -> bool:
        """Return True if restore completed successfully."""
        return self.status == CheckpointStatus.COMPLETED

    @property
    def duration_seconds(self) -> float | None:
        """Calculate duration of restore operation."""
        if self.completed_at is None:
            return None
        return (self.completed_at - self.started_at).total_seconds()


class MigrationResult(BaseModel):
    """Result of a live migration operation."""

    model_config = ConfigDict(frozen=False)

    container_name: str
    source_host: str
    target_host: str
    checkpoint_name: str
    checkpoint_id: str

    started_at: datetime = Field(default_factory=datetime.now)
    completed_at: datetime | None = None
    status: CheckpointStatus = CheckpointStatus.PENDING

    # Transfer details
    checkpoint_size_bytes: int = 0
    transfer_size_bytes: int = 0
    transfer_duration_seconds: float | None = None

    # New container on target
    target_container_id: str | None = None

    error_message: str | None = None
    error_output: str | None = None

    @property
    def success(self) -> bool:
        """Return True if migration completed successfully."""
        return self.status == CheckpointStatus.COMPLETED

    @property
    def duration_seconds(self) -> float | None:
        """Calculate total duration of migration."""
        if self.completed_at is None:
            return None
        return (self.completed_at - self.started_at).total_seconds()


class CheckpointPolicy(BaseModel):
    """Checkpoint policy for a container."""

    model_config = ConfigDict(frozen=True)

    container_name: Annotated[str, Field(min_length=1, max_length=64)] = Field(
        ..., description="Container name"
    )

    # Checkpoint options
    checkpoint_type: CheckpointType = Field(
        default=CheckpointType.FULL, description="Default checkpoint type"
    )
    leave_running: bool = Field(
        default=False, description="Keep container running after checkpoint"
    )
    tcp_established: bool = Field(
        default=False, description="Checkpoint established TCP connections"
    )
    external_unix_sockets: bool = Field(
        default=False, description="Allow external unix sockets"
    )

    # Retention
    max_checkpoints: Annotated[int | None, Field(ge=1)] = Field(
        default=10, description="Maximum number of checkpoints to keep"
    )
    max_age_days: Annotated[int | None, Field(ge=1)] = Field(
        default=30, description="Delete checkpoints older than N days"
    )
    max_total_size_mb: Annotated[int | None, Field(ge=1)] = Field(
        default=None, description="Max total checkpoint size in MB"
    )

    # Pre/post commands
    pre_checkpoint_commands: list[str] = Field(
        default_factory=list, description="Commands to run before checkpoint"
    )
    post_checkpoint_commands: list[str] = Field(
        default_factory=list, description="Commands to run after checkpoint"
    )

    verify_after_checkpoint: bool = Field(
        default=True, description="Verify checkpoint integrity after creation"
    )


class CheckpointHistory(BaseModel):
    """History of checkpoint operations for a container."""

    model_config = ConfigDict(frozen=False)

    container_name: str
    checkpoints: list[CheckpointMetadata] = Field(default_factory=list)
    last_checkpoint: datetime | None = None
    last_successful_checkpoint: datetime | None = None
    total_checkpoints: int = 0
    successful_checkpoints: int = 0
    failed_checkpoints: int = 0
    total_size_bytes: int = 0

    def add_checkpoint(self, metadata: CheckpointMetadata) -> None:
        """Add a checkpoint to history."""
        self.checkpoints.append(metadata)
        self.last_checkpoint = metadata.created_at
        self.total_checkpoints += 1

        if metadata.completed_at:
            self.successful_checkpoints += 1
            self.last_successful_checkpoint = metadata.completed_at
            self.total_size_bytes += metadata.size_bytes
        else:
            self.failed_checkpoints += 1

        # Keep only last 100 entries in memory (checkpoints are larger than backups)
        if len(self.checkpoints) > 100:
            self.checkpoints = self.checkpoints[-100:]

    def remove_checkpoint(self, checkpoint_id: str) -> bool:
        """Remove a checkpoint from history by ID."""
        for i, checkpoint in enumerate(self.checkpoints):
            if checkpoint.checkpoint_id == checkpoint_id:
                removed = self.checkpoints.pop(i)
                self.total_size_bytes -= removed.size_bytes
                return True
        return False

    def get_checkpoint(self, checkpoint_name: str) -> CheckpointMetadata | None:
        """Get a checkpoint by name."""
        for checkpoint in self.checkpoints:
            if checkpoint.checkpoint_name == checkpoint_name:
                return checkpoint
        return None

    def get_checkpoint_by_id(self, checkpoint_id: str) -> CheckpointMetadata | None:
        """Get a checkpoint by ID."""
        for checkpoint in self.checkpoints:
            if checkpoint.checkpoint_id == checkpoint_id:
                return checkpoint
        return None


class CRIUCapabilities(BaseModel):
    """CRIU feature detection results."""

    model_config = ConfigDict(frozen=True)

    available: bool = Field(
        default=False, description="Whether CRIU is installed and working"
    )
    version: str | None = Field(default=None, description="CRIU version string")

    # Kernel features
    kernel_checkpoint_restore: bool = Field(
        default=False, description="CONFIG_CHECKPOINT_RESTORE enabled"
    )
    kernel_ns_last_pid: bool = Field(
        default=False, description="/proc/sys/kernel/ns_last_pid exists"
    )
    kernel_userfaultfd: bool = Field(
        default=False, description="userfaultfd support for lazy migration"
    )

    # CRIU features
    tcp_established: bool = Field(
        default=False, description="Can checkpoint TCP connections"
    )
    external_unix_sockets: bool = Field(
        default=False, description="Can checkpoint external unix sockets"
    )
    file_locks: bool = Field(default=False, description="Can checkpoint file locks")
    lazy_pages: bool = Field(
        default=False, description="Lazy page transfer for migration"
    )

    # Permissions
    has_cap_sys_admin: bool = Field(
        default=False, description="CAP_SYS_ADMIN capability available"
    )
    can_checkpoint_namespaces: bool = Field(
        default=False, description="Can checkpoint namespaced processes"
    )


class DockerCheckpointCapabilities(BaseModel):
    """Docker checkpoint feature detection results."""

    model_config = ConfigDict(frozen=True)

    available: bool = Field(
        default=False, description="Whether Docker is available"
    )
    experimental_enabled: bool = Field(
        default=False, description="Docker experimental mode enabled"
    )
    checkpoint_supported: bool = Field(
        default=False, description="Docker checkpoint command available"
    )
    version: str | None = Field(default=None, description="Docker version string")
