"""Pydantic models for security domain-inspired security patterns.

This module defines data structures for:
- Template containers (base images for spawning app containers)
- Disposable containers (ephemeral, destroyed after use)
- Security domains (color-coded trust levels)
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Annotated

from pydantic import BaseModel, ConfigDict, Field

# =============================================================================
# Template Container Models
# =============================================================================


class TemplateType(str, Enum):
    """Type of template container."""

    SYSTEM = "system"  # Base system template (minimal, for other templates)
    APP = "app"  # Application template (with GUI support)
    MINIMAL = "minimal"  # Minimal template (bare essentials)
    NETWORK = "network"  # Network service template (VPN, firewall)


class TemplateConfig(BaseModel):
    """Configuration for creating a template container.

    Templates are read-only base images that spawn app containers.
    App containers use OverlayFS to layer changes on top of templates.

    Example:
        >>> config = TemplateConfig(
        ...     name="debian-dev",
        ...     distro="trixie",
        ...     profiles=["base", "dev-python"],
        ...     template_type=TemplateType.APP,
        ... )
    """

    model_config = ConfigDict(frozen=True)

    name: Annotated[
        str,
        Field(
            min_length=1,
            max_length=64,
            pattern=r"^[a-zA-Z0-9][a-zA-Z0-9_.-]*$",
            description="Unique template name",
        ),
    ]

    distro: Annotated[str, Field(default="trixie", description="Base distribution")]

    profiles: Annotated[
        list[str],
        Field(default_factory=list, description="Profiles to apply to template"),
    ]

    template_type: Annotated[
        TemplateType,
        Field(default=TemplateType.APP, description="Template type"),
    ]

    packages: Annotated[
        list[str],
        Field(default_factory=list, description="Additional packages to install"),
    ]

    description: Annotated[str | None, Field(default=None, description="Template description")]

    # Resource limits for spawned containers (defaults)
    default_memory_mb: Annotated[
        int | None, Field(default=None, ge=128, description="Default memory limit for instances")
    ]
    default_cpu_count: Annotated[
        float | None, Field(default=None, ge=0.5, description="Default CPU limit for instances")
    ]


class TemplateState(BaseModel):
    """Runtime state of a template container.

    Tracks the template's filesystem location, version, and spawned instances.
    """

    model_config = ConfigDict(frozen=False)

    name: str
    """Template name."""

    distro: str
    """Base distribution."""

    template_type: TemplateType
    """Template type."""

    profiles: list[str] = Field(default_factory=list)
    """Applied profiles."""

    root_path: Path
    """Path to template root filesystem."""

    version: int = Field(default=1, ge=1)
    """Template version (incremented on updates)."""

    created_at: datetime = Field(default_factory=datetime.now)
    """When the template was created."""

    updated_at: datetime | None = None
    """When the template was last updated."""

    instances: list[str] = Field(default_factory=list)
    """Names of containers spawned from this template."""

    size_bytes: int = Field(default=0, ge=0)
    """Size of template filesystem in bytes."""

    description: str | None = None
    """Template description."""

    default_memory_mb: int | None = None
    """Default memory limit for instances."""

    default_cpu_count: float | None = None
    """Default CPU limit for instances."""


class TemplateInstance(BaseModel):
    """A container instance spawned from a template.

    Instances use OverlayFS to layer changes on top of the template.
    Ephemeral instances are destroyed on stop; persistent instances keep changes.
    """

    model_config = ConfigDict(frozen=False)

    name: str
    """Instance name."""

    template_name: str
    """Name of the source template."""

    template_version: int
    """Version of the template when spawned."""

    ephemeral: bool = False
    """If True, instance is destroyed on stop."""

    overlay_path: Path | None = None
    """Path to overlay upperdir (for persistent instances)."""

    work_path: Path | None = None
    """Path to overlay workdir."""

    created_at: datetime = Field(default_factory=datetime.now)
    """When the instance was created."""

    last_started: datetime | None = None
    """When the instance was last started."""

    preserve_paths: list[str] = Field(default_factory=list)
    """Paths to preserve even for ephemeral instances (e.g., ~/Downloads)."""


# =============================================================================
# Disposable Container Models
# =============================================================================


class DisposableType(str, Enum):
    """Pre-configured disposable container types."""

    BROWSER = "browser"  # Web browser for untrusted browsing
    FILE_VIEWER = "file_viewer"  # Document/media viewer
    TERMINAL = "terminal"  # Terminal for untrusted commands
    CUSTOM = "custom"  # Custom application


class DisposableConfig(BaseModel):
    """Configuration for a disposable container type.

    Disposables are ephemeral containers spawned from templates.
    Each launch creates a fresh instance that is destroyed after use.

    Example:
        >>> config = DisposableConfig(
        ...     name="browser",
        ...     template_name="debian-browser",
        ...     disposable_type=DisposableType.BROWSER,
        ...     application="/usr/bin/firefox",
        ...     preserve_downloads=True,
        ... )
    """

    model_config = ConfigDict(frozen=True)

    name: Annotated[
        str,
        Field(
            min_length=1,
            max_length=64,
            pattern=r"^[a-zA-Z0-9][a-zA-Z0-9_.-]*$",
            description="Unique disposable type name",
        ),
    ]

    template_name: Annotated[str, Field(description="Template to spawn from")]

    disposable_type: Annotated[
        DisposableType,
        Field(default=DisposableType.CUSTOM, description="Disposable type"),
    ]

    application: Annotated[
        str | None,
        Field(default=None, description="Application to launch (e.g., /usr/bin/firefox)"),
    ]

    application_args: Annotated[
        list[str],
        Field(default_factory=list, description="Default application arguments"),
    ]

    description: Annotated[str | None, Field(default=None, description="Description")]

    # Preserve specific paths between launches
    preserve_downloads: Annotated[
        bool, Field(default=False, description="Preserve ~/Downloads between launches")
    ]
    preserve_paths: Annotated[
        list[str],
        Field(default_factory=list, description="Additional paths to preserve"),
    ]

    # Desktop integration
    icon: Annotated[str | None, Field(default=None, description="Icon name or path")]
    icon_color: Annotated[
        str | None, Field(default=None, description="Hex color to tint icon")
    ]
    desktop_name: Annotated[
        str | None, Field(default=None, description="Name for .desktop file")
    ]
    desktop_categories: Annotated[
        list[str], Field(default_factory=list, description="XDG categories")
    ]
    mime_types: Annotated[
        list[str],
        Field(default_factory=list, description="MIME types for file associations"),
    ]

    # Security settings
    domain: Annotated[
        str | None, Field(default=None, description="Security domain (e.g., 'untrusted')")
    ]
    allow_network: Annotated[bool, Field(default=True, description="Allow network access")]


class DisposableInstance(BaseModel):
    """A running disposable container instance.

    Tracks a single launch of a disposable container.
    Destroyed automatically when the application exits.
    """

    model_config = ConfigDict(frozen=False)

    instance_id: str
    """Unique instance ID (e.g., 'disp-browser-1234')."""

    disposable_name: str
    """Name of the disposable configuration."""

    template_name: str
    """Template it was spawned from."""

    container_name: str
    """Actual container name in the runtime."""

    pid: int | None = None
    """PID of the main application process."""

    started_at: datetime = Field(default_factory=datetime.now)
    """When the instance was started."""

    arguments: list[str] = Field(default_factory=list)
    """Arguments passed to the application."""


# =============================================================================
# Security Domain Models
# =============================================================================


class SecurityDomain(str, Enum):
    """Pre-defined security domains (trust levels).

    Based on security domain security domains, adapted for containers.
    """

    VAULT = "vault"  # Highest trust, no network, stores secrets
    WORK = "work"  # Trusted work environment
    PERSONAL = "personal"  # Personal use, moderate trust
    SELF_HOSTED = "self-hosted"  # Self-hosted services (git, wiki, etc.)
    MONITORING = "monitoring"  # Observability infrastructure (metrics, logs, alerts)
    UNTRUSTED = "untrusted"  # Low trust, heavy restrictions
    DISPOSABLE = "disposable"  # Ephemeral, very restricted


# Default colors for each domain (matching security domain conventions)
DOMAIN_COLORS: dict[SecurityDomain, str] = {
    SecurityDomain.VAULT: "#800080",  # Purple
    SecurityDomain.WORK: "#0066B3",  # Blue
    SecurityDomain.PERSONAL: "#00AA00",  # Green
    SecurityDomain.SELF_HOSTED: "#FF8C00",  # Orange
    SecurityDomain.MONITORING: "#008080",  # Teal
    SecurityDomain.UNTRUSTED: "#CC0000",  # Red
    SecurityDomain.DISPOSABLE: "#808080",  # Gray
}


class DomainPolicy(BaseModel):
    """Security policy for a domain.

    Defines what containers in this domain can do.

    Example:
        >>> policy = DomainPolicy(
        ...     domain=SecurityDomain.VAULT,
        ...     color="#800080",
        ...     allow_network=False,
        ...     ephemeral_only=False,
        ...     required_profiles=["security-hardened"],
        ... )
    """

    model_config = ConfigDict(frozen=True)

    domain: SecurityDomain
    """The security domain."""

    name: Annotated[str | None, Field(default=None, description="Custom domain name")]

    color: Annotated[str, Field(pattern=r"^#[0-9A-Fa-f]{6}$", description="Hex color")]

    description: Annotated[str | None, Field(default=None)]

    # Network policy
    allow_network: Annotated[bool, Field(default=True, description="Allow network access")]
    allow_internet: Annotated[
        bool, Field(default=True, description="Allow internet access (vs local only)")
    ]
    required_gateway: Annotated[
        str | None, Field(default=None, description="Required gateway container (VPN/Tor)")
    ]

    # Container policy
    ephemeral_only: Annotated[
        bool, Field(default=False, description="Only allow ephemeral containers")
    ]
    read_only_root: Annotated[
        bool, Field(default=False, description="Force read-only root filesystem")
    ]

    # Resource limits
    max_memory_mb: Annotated[
        int | None, Field(default=None, ge=128, description="Maximum memory per container")
    ]
    max_cpu_percent: Annotated[
        int | None, Field(default=None, ge=1, le=100, description="Maximum CPU percentage")
    ]

    # Required profiles for containers in this domain
    required_profiles: Annotated[
        list[str],
        Field(default_factory=list, description="Profiles automatically applied"),
    ]

    # Inter-domain communication
    allow_communication_with: Annotated[
        list[SecurityDomain],
        Field(default_factory=list, description="Domains this can communicate with"),
    ]


# Default policies for each domain
DEFAULT_DOMAIN_POLICIES: dict[SecurityDomain, DomainPolicy] = {
    SecurityDomain.VAULT: DomainPolicy(
        domain=SecurityDomain.VAULT,
        color=DOMAIN_COLORS[SecurityDomain.VAULT],
        description="Secure vault for secrets and keys",
        allow_network=False,
        allow_internet=False,
        ephemeral_only=False,
        read_only_root=False,
        required_profiles=["security-hardened"],
        allow_communication_with=[],
    ),
    SecurityDomain.WORK: DomainPolicy(
        domain=SecurityDomain.WORK,
        color=DOMAIN_COLORS[SecurityDomain.WORK],
        description="Trusted work environment",
        allow_network=True,
        allow_internet=True,
        ephemeral_only=False,
        read_only_root=False,
        required_profiles=[],
        allow_communication_with=[SecurityDomain.PERSONAL],
    ),
    SecurityDomain.PERSONAL: DomainPolicy(
        domain=SecurityDomain.PERSONAL,
        color=DOMAIN_COLORS[SecurityDomain.PERSONAL],
        description="Personal use containers",
        allow_network=True,
        allow_internet=True,
        ephemeral_only=False,
        read_only_root=False,
        required_profiles=[],
        allow_communication_with=[SecurityDomain.WORK],
    ),
    SecurityDomain.SELF_HOSTED: DomainPolicy(
        domain=SecurityDomain.SELF_HOSTED,
        color=DOMAIN_COLORS[SecurityDomain.SELF_HOSTED],
        description="Self-hosted services (git forges, wikis, media servers)",
        allow_network=True,
        allow_internet=True,
        ephemeral_only=False,
        read_only_root=False,
        required_profiles=[],
        allow_communication_with=[SecurityDomain.PERSONAL, SecurityDomain.WORK],
    ),
    SecurityDomain.MONITORING: DomainPolicy(
        domain=SecurityDomain.MONITORING,
        color=DOMAIN_COLORS[SecurityDomain.MONITORING],
        description="Observability infrastructure (metrics, logs, alerts, notifications)",
        allow_network=True,
        allow_internet=True,
        ephemeral_only=False,
        read_only_root=False,
        required_profiles=[],
        allow_communication_with=[SecurityDomain.SELF_HOSTED, SecurityDomain.WORK],
    ),
    SecurityDomain.UNTRUSTED: DomainPolicy(
        domain=SecurityDomain.UNTRUSTED,
        color=DOMAIN_COLORS[SecurityDomain.UNTRUSTED],
        description="Untrusted content handling",
        allow_network=True,
        allow_internet=True,
        ephemeral_only=False,
        read_only_root=True,
        max_memory_mb=2048,
        max_cpu_percent=50,
        required_profiles=["security-hardened"],
        allow_communication_with=[],
    ),
    SecurityDomain.DISPOSABLE: DomainPolicy(
        domain=SecurityDomain.DISPOSABLE,
        color=DOMAIN_COLORS[SecurityDomain.DISPOSABLE],
        description="Ephemeral disposable containers",
        allow_network=True,
        allow_internet=True,
        ephemeral_only=True,
        read_only_root=True,
        max_memory_mb=2048,
        max_cpu_percent=50,
        required_profiles=[],
        allow_communication_with=[],
    ),
}


# =============================================================================
# Service Container Models (sys-tor, sys-vpn, sys-firewall)
# =============================================================================


class ServiceContainerType(str, Enum):
    """Type of service container (like sys- VMs)."""

    TOR_GATEWAY = "tor_gateway"  # sys-whonix equivalent - Tor routing
    VPN_GATEWAY = "vpn_gateway"  # VPN gateway (OpenVPN, WireGuard)
    FIREWALL = "firewall"  # sys-firewall equivalent - filtering/routing
    DNS = "dns"  # DNS resolver/filter
    PROXY = "proxy"  # HTTP/SOCKS proxy


class ServiceContainerConfig(BaseModel):
    """Configuration for creating a service container.

    Service containers provide network services to other containers,
    similar to security domain sys-net, sys-firewall, sys-whonix.

    Example:
        >>> config = ServiceContainerConfig(
        ...     name="sys-tor",
        ...     service_type=ServiceContainerType.TOR_GATEWAY,
        ...     transparent_proxy=True,
        ... )
    """

    model_config = ConfigDict(frozen=True)

    name: Annotated[
        str,
        Field(
            min_length=1,
            max_length=64,
            pattern=r"^[a-zA-Z0-9][a-zA-Z0-9_.-]*$",
            description="Unique service container name",
        ),
    ]

    service_type: Annotated[
        ServiceContainerType,
        Field(description="Type of service"),
    ]

    distro: Annotated[str, Field(default="trixie", description="Base distribution")]

    description: Annotated[str | None, Field(default=None)]

    # Tor-specific options
    transparent_proxy: Annotated[
        bool, Field(default=True, description="Use transparent proxy (vs SOCKS)")
    ]
    tor_exit_nodes: Annotated[
        list[str],
        Field(default_factory=list, description="Preferred exit node countries"),
    ]
    tor_exclude_nodes: Annotated[
        list[str],
        Field(default_factory=list, description="Countries to exclude"),
    ]
    tor_bridges: Annotated[
        list[str],
        Field(default_factory=list, description="Bridge addresses for censorship"),
    ]

    # VPN-specific options
    vpn_config_file: Annotated[
        str | None,
        Field(default=None, description="Path to VPN config (OpenVPN/WireGuard)"),
    ]
    vpn_provider: Annotated[
        str | None,
        Field(default=None, description="VPN provider type: openvpn or wireguard"),
    ]

    # Common options
    kill_switch: Annotated[
        bool, Field(default=True, description="Block traffic if service drops")
    ]
    dns_servers: Annotated[
        list[str],
        Field(default_factory=list, description="DNS servers to use"),
    ]
    dns_leak_prevention: Annotated[
        bool, Field(default=True, description="Prevent DNS leaks")
    ]

    # Firewall-specific options
    allow_ports: Annotated[
        list[int],
        Field(default_factory=list, description="Ports to allow through firewall"),
    ]
    block_ipv6: Annotated[
        bool, Field(default=True, description="Block all IPv6 traffic")
    ]


class ServiceContainerState(BaseModel):
    """Runtime state of a service container."""

    model_config = ConfigDict(frozen=False)

    name: str
    """Service container name."""

    service_type: ServiceContainerType
    """Type of service."""

    container_name: str
    """Actual nspawn container name."""

    status: str = "stopped"
    """Current status: stopped, starting, running, error."""

    created_at: datetime = Field(default_factory=datetime.now)
    """When the service was created."""

    started_at: datetime | None = None
    """When the service was last started."""

    # Service-specific state
    public_ip: str | None = None
    """Public IP address (for VPN/Tor)."""

    tor_bootstrap_progress: int = 0
    """Tor bootstrap progress (0-100)."""

    vpn_interface: str | None = None
    """VPN interface name (tun0, wg0)."""

    # Client tracking
    attached_clients: list[str] = Field(default_factory=list)
    """Containers routing traffic through this service."""

    error_message: str | None = None
    """Last error message if status is 'error'."""


class ClientAttachment(BaseModel):
    """Tracks a client container attached to a service container."""

    model_config = ConfigDict(frozen=False)

    client_name: str
    """Name of the client container."""

    service_name: str
    """Name of the service container."""

    attached_at: datetime = Field(default_factory=datetime.now)
    """When the attachment was made."""

    routing_mode: str = "all"
    """Routing mode: 'all' (all traffic) or 'selected' (specific ports)."""

    selected_ports: list[int] = Field(default_factory=list)
    """Ports to route (if routing_mode is 'selected')."""


# =============================================================================
# Store Models
# =============================================================================


class ServiceContainerStore(BaseModel):
    """Persisted service container registry."""

    model_config = ConfigDict(frozen=False)

    version: int = 1
    services: dict[str, ServiceContainerState] = Field(default_factory=dict)
    attachments: dict[str, ClientAttachment] = Field(default_factory=dict)

    def add_service(self, service: ServiceContainerState) -> None:
        """Add or update a service container."""
        self.services[service.name] = service

    def get_service(self, name: str) -> ServiceContainerState | None:
        """Get a service by name."""
        return self.services.get(name)

    def remove_service(self, name: str) -> bool:
        """Remove a service. Returns True if removed."""
        if name in self.services:
            del self.services[name]
            return True
        return False

    def list_services(
        self, service_type: ServiceContainerType | None = None
    ) -> list[ServiceContainerState]:
        """List all services, optionally filtered by type."""
        services = list(self.services.values())
        if service_type:
            services = [s for s in services if s.service_type == service_type]
        return services

    def add_attachment(self, attachment: ClientAttachment) -> None:
        """Track a client attachment."""
        key = f"{attachment.client_name}:{attachment.service_name}"
        self.attachments[key] = attachment

    def remove_attachment(self, client_name: str, service_name: str) -> bool:
        """Remove an attachment."""
        key = f"{client_name}:{service_name}"
        if key in self.attachments:
            del self.attachments[key]
            return True
        return False

    def get_client_attachments(self, client_name: str) -> list[ClientAttachment]:
        """Get all services a client is attached to."""
        return [a for a in self.attachments.values() if a.client_name == client_name]

    def get_service_clients(self, service_name: str) -> list[ClientAttachment]:
        """Get all clients attached to a service."""
        return [a for a in self.attachments.values() if a.service_name == service_name]


class TemplateStore(BaseModel):
    """Persisted template registry."""

    model_config = ConfigDict(frozen=False)

    version: int = 1
    templates: dict[str, TemplateState] = Field(default_factory=dict)

    def add_template(self, template: TemplateState) -> None:
        """Add or update a template."""
        self.templates[template.name] = template

    def get_template(self, name: str) -> TemplateState | None:
        """Get a template by name."""
        return self.templates.get(name)

    def remove_template(self, name: str) -> bool:
        """Remove a template. Returns True if removed."""
        if name in self.templates:
            del self.templates[name]
            return True
        return False

    def list_templates(
        self, template_type: TemplateType | None = None
    ) -> list[TemplateState]:
        """List all templates, optionally filtered by type."""
        templates = list(self.templates.values())
        if template_type:
            templates = [t for t in templates if t.template_type == template_type]
        return templates


class DisposableStore(BaseModel):
    """Persisted disposable configuration registry."""

    model_config = ConfigDict(frozen=False)

    version: int = 1
    disposables: dict[str, DisposableConfig] = Field(default_factory=dict)
    instances: dict[str, DisposableInstance] = Field(default_factory=dict)

    def add_disposable(self, config: DisposableConfig) -> None:
        """Add or update a disposable configuration."""
        # Convert frozen model to dict and back to store
        self.disposables[config.name] = config

    def get_disposable(self, name: str) -> DisposableConfig | None:
        """Get a disposable configuration by name."""
        return self.disposables.get(name)

    def remove_disposable(self, name: str) -> bool:
        """Remove a disposable configuration. Returns True if removed."""
        if name in self.disposables:
            del self.disposables[name]
            return True
        return False

    def list_disposables(
        self, disposable_type: DisposableType | None = None
    ) -> list[DisposableConfig]:
        """List all disposable configs, optionally filtered by type."""
        configs = list(self.disposables.values())
        if disposable_type:
            configs = [c for c in configs if c.disposable_type == disposable_type]
        return configs

    def add_instance(self, instance: DisposableInstance) -> None:
        """Track a running instance."""
        self.instances[instance.instance_id] = instance

    def remove_instance(self, instance_id: str) -> bool:
        """Remove an instance record. Returns True if removed."""
        if instance_id in self.instances:
            del self.instances[instance_id]
            return True
        return False

    def list_instances(self, disposable_name: str | None = None) -> list[DisposableInstance]:
        """List running instances, optionally filtered by disposable name."""
        instances = list(self.instances.values())
        if disposable_name:
            instances = [i for i in instances if i.disposable_name == disposable_name]
        return instances


class DomainStore(BaseModel):
    """Persisted domain policy registry."""

    model_config = ConfigDict(frozen=False)

    version: int = 1
    custom_policies: dict[str, DomainPolicy] = Field(default_factory=dict)
    container_domains: dict[str, SecurityDomain] = Field(default_factory=dict)

    def get_policy(self, domain: SecurityDomain) -> DomainPolicy:
        """Get policy for a domain (custom or default)."""
        if domain.value in self.custom_policies:
            return self.custom_policies[domain.value]
        return DEFAULT_DOMAIN_POLICIES[domain]

    def set_custom_policy(self, policy: DomainPolicy) -> None:
        """Set a custom policy for a domain."""
        self.custom_policies[policy.domain.value] = policy

    def assign_container(self, container_name: str, domain: SecurityDomain) -> None:
        """Assign a container to a domain."""
        self.container_domains[container_name] = domain

    def get_container_domain(self, container_name: str) -> SecurityDomain | None:
        """Get the domain assigned to a container."""
        return self.container_domains.get(container_name)

    def unassign_container(self, container_name: str) -> bool:
        """Remove a container's domain assignment."""
        if container_name in self.container_domains:
            del self.container_domains[container_name]
            return True
        return False

    def list_containers_in_domain(self, domain: SecurityDomain) -> list[str]:
        """List all containers assigned to a domain."""
        return [
            name for name, d in self.container_domains.items() if d == domain
        ]


# =============================================================================
# Split Secrets Vault Models
# =============================================================================


class VaultType(str, Enum):
    """Type of secrets vault."""

    GPG = "gpg"  # GPG keys and signing
    SSH = "ssh"  # SSH keys and authentication
    COMBINED = "combined"  # Both GPG and SSH


class VaultOperation(str, Enum):
    """Operations that can be performed via a vault."""

    # GPG operations
    SIGN = "sign"  # Sign data
    DECRYPT = "decrypt"  # Decrypt data
    ENCRYPT = "encrypt"  # Encrypt data (public key, but vault may hold it)

    # SSH operations
    SSH_AUTH = "ssh_auth"  # SSH authentication
    SSH_SIGN = "ssh_sign"  # SSH signing (e.g., git commits)

    # Administrative
    EXPORT_PUBLIC = "export_public"  # Export public keys
    LIST_KEYS = "list_keys"  # List available keys


class VaultKeyInfo(BaseModel):
    """Information about a key stored in a vault."""

    model_config = ConfigDict(frozen=False)

    key_id: str
    """Key identifier (GPG fingerprint or SSH key fingerprint)."""

    key_type: str
    """Key type (e.g., 'rsa4096', 'ed25519', 'cv25519')."""

    vault_type: VaultType
    """Whether this is a GPG or SSH key."""

    user_id: str | None = None
    """User ID for GPG keys (e.g., 'User Name <email@example.com>')."""

    created_at: datetime | None = None
    """When the key was created."""

    imported_at: datetime = Field(default_factory=datetime.now)
    """When the key was imported into the vault."""

    comment: str | None = None
    """Optional comment about the key."""

    # For SSH keys
    ssh_public_key: str | None = None
    """SSH public key (for authorized_keys)."""


class VaultConfig(BaseModel):
    """Configuration for creating a vault container.

    Vaults are containers with NO network access that store private keys.
    Other containers can request cryptographic operations without
    ever seeing the private keys.

    Example:
        >>> config = VaultConfig(
        ...     name="gpg-vault",
        ...     vault_type=VaultType.GPG,
        ...     description="GPG signing keys",
        ... )
    """

    model_config = ConfigDict(frozen=True)

    name: Annotated[
        str,
        Field(
            min_length=1,
            max_length=64,
            pattern=r"^[a-zA-Z0-9][a-zA-Z0-9_.-]*$",
            description="Unique vault name",
        ),
    ]

    vault_type: Annotated[
        VaultType,
        Field(description="Type of vault (gpg, ssh, or combined)"),
    ]

    distro: Annotated[str, Field(default="trixie", description="Base distribution")]

    description: Annotated[str | None, Field(default=None)]

    # Security settings
    auto_lock_minutes: Annotated[
        int | None,
        Field(default=30, ge=1, description="Auto-lock after inactivity (minutes)"),
    ]

    require_confirmation: Annotated[
        bool,
        Field(default=True, description="Require user confirmation for operations"),
    ]

    log_operations: Annotated[
        bool,
        Field(default=True, description="Log all cryptographic operations"),
    ]


class VaultPolicy(BaseModel):
    """Policy controlling what operations a container can request.

    Example:
        >>> policy = VaultPolicy(
        ...     container_name="work-email",
        ...     vault_name="gpg-vault",
        ...     allowed_operations=[VaultOperation.SIGN],
        ...     key_filter=["0x1234567890ABCDEF"],  # Only specific key
        ... )
    """

    model_config = ConfigDict(frozen=True)

    container_name: Annotated[str, Field(description="Container allowed to use vault")]

    vault_name: Annotated[str, Field(description="Vault this policy applies to")]

    allowed_operations: Annotated[
        list[VaultOperation],
        Field(description="Operations this container can perform"),
    ]

    key_filter: Annotated[
        list[str],
        Field(
            default_factory=list,
            description="Key IDs this policy applies to (empty = all keys)",
        ),
    ]

    require_confirmation: Annotated[
        bool | None,
        Field(default=None, description="Override vault's require_confirmation setting"),
    ]

    expires_at: Annotated[
        datetime | None,
        Field(default=None, description="When this policy expires"),
    ]

    description: Annotated[str | None, Field(default=None)]


class VaultAttachment(BaseModel):
    """Tracks a container attached to a vault for agent forwarding."""

    model_config = ConfigDict(frozen=False)

    container_name: str
    """Name of the attached container."""

    vault_name: str
    """Name of the vault."""

    attached_at: datetime = Field(default_factory=datetime.now)
    """When the attachment was made."""

    socket_path: str | None = None
    """Path to the forwarded agent socket in the container."""

    gpg_socket_path: str | None = None
    """Path to the GPG agent socket (if GPG vault)."""

    ssh_socket_path: str | None = None
    """Path to the SSH agent socket (if SSH vault)."""


class VaultState(BaseModel):
    """Runtime state of a vault container."""

    model_config = ConfigDict(frozen=False)

    name: str
    """Vault name."""

    vault_type: VaultType
    """Type of vault."""

    container_name: str
    """Actual nspawn container name."""

    status: str = "stopped"
    """Current status: stopped, starting, running, locked, error."""

    created_at: datetime = Field(default_factory=datetime.now)
    """When the vault was created."""

    started_at: datetime | None = None
    """When the vault was last started."""

    locked_at: datetime | None = None
    """When the vault was last locked (if locked)."""

    # Key inventory
    keys: list[VaultKeyInfo] = Field(default_factory=list)
    """Keys stored in this vault."""

    # Attached containers
    attached_containers: list[str] = Field(default_factory=list)
    """Containers currently attached to this vault."""

    # Agent socket paths (inside the vault container)
    gpg_agent_socket: str | None = None
    """Path to GPG agent socket."""

    ssh_agent_socket: str | None = None
    """Path to SSH agent socket."""

    error_message: str | None = None
    """Last error message if status is 'error'."""

    description: str | None = None
    """Vault description."""


class VaultStore(BaseModel):
    """Persisted vault registry."""

    model_config = ConfigDict(frozen=False)

    version: int = 1
    vaults: dict[str, VaultState] = Field(default_factory=dict)
    policies: dict[str, VaultPolicy] = Field(default_factory=dict)
    attachments: dict[str, VaultAttachment] = Field(default_factory=dict)

    def add_vault(self, vault: VaultState) -> None:
        """Add or update a vault."""
        self.vaults[vault.name] = vault

    def get_vault(self, name: str) -> VaultState | None:
        """Get a vault by name."""
        return self.vaults.get(name)

    def remove_vault(self, name: str) -> bool:
        """Remove a vault. Returns True if removed."""
        if name in self.vaults:
            del self.vaults[name]
            # Also remove associated policies and attachments
            self.policies = {
                k: v for k, v in self.policies.items() if v.vault_name != name
            }
            self.attachments = {
                k: v for k, v in self.attachments.items() if v.vault_name != name
            }
            return True
        return False

    def list_vaults(self, vault_type: VaultType | None = None) -> list[VaultState]:
        """List all vaults, optionally filtered by type."""
        vaults = list(self.vaults.values())
        if vault_type:
            vaults = [v for v in vaults if v.vault_type == vault_type]
        return vaults

    def add_policy(self, policy: VaultPolicy) -> None:
        """Add or update a policy."""
        key = f"{policy.container_name}:{policy.vault_name}"
        self.policies[key] = policy

    def get_policy(self, container_name: str, vault_name: str) -> VaultPolicy | None:
        """Get policy for a container-vault pair."""
        key = f"{container_name}:{vault_name}"
        return self.policies.get(key)

    def remove_policy(self, container_name: str, vault_name: str) -> bool:
        """Remove a policy."""
        key = f"{container_name}:{vault_name}"
        if key in self.policies:
            del self.policies[key]
            return True
        return False

    def list_policies(self, vault_name: str | None = None) -> list[VaultPolicy]:
        """List policies, optionally filtered by vault."""
        policies = list(self.policies.values())
        if vault_name:
            policies = [p for p in policies if p.vault_name == vault_name]
        return policies

    def get_container_policies(self, container_name: str) -> list[VaultPolicy]:
        """Get all policies for a specific container."""
        return [p for p in self.policies.values() if p.container_name == container_name]

    def add_attachment(self, attachment: VaultAttachment) -> None:
        """Track a container attachment."""
        key = f"{attachment.container_name}:{attachment.vault_name}"
        self.attachments[key] = attachment

    def remove_attachment(self, container_name: str, vault_name: str) -> bool:
        """Remove an attachment."""
        key = f"{container_name}:{vault_name}"
        if key in self.attachments:
            del self.attachments[key]
            return True
        return False

    def get_attachment(
        self, container_name: str, vault_name: str
    ) -> VaultAttachment | None:
        """Get attachment for a container-vault pair."""
        key = f"{container_name}:{vault_name}"
        return self.attachments.get(key)

    def list_attachments(self, vault_name: str | None = None) -> list[VaultAttachment]:
        """List attachments, optionally filtered by vault."""
        attachments = list(self.attachments.values())
        if vault_name:
            attachments = [a for a in attachments if a.vault_name == vault_name]
        return attachments


# =============================================================================
# Network Chain Models
# =============================================================================


class ChainHopType(str, Enum):
    """Type of hop in a network chain."""

    APP = "app"  # Application container (source of traffic)
    FIREWALL = "firewall"  # Firewall/filter container
    GATEWAY = "gateway"  # Gateway container (tor, vpn)
    PROXY = "proxy"  # Proxy container


class ChainHop(BaseModel):
    """A single hop in a network chain.

    Example:
        >>> hop = ChainHop(
        ...     hop_type=ChainHopType.GATEWAY,
        ...     container_name="sys-tor",
        ...     position=2,
        ... )
    """

    model_config = ConfigDict(frozen=True)

    hop_type: Annotated[ChainHopType, Field(description="Type of this hop")]

    container_name: Annotated[str, Field(description="Container providing this hop")]

    position: Annotated[int, Field(ge=0, description="Position in chain (0 = app)")]

    # Policy at this hop
    log_traffic: Annotated[
        bool, Field(default=False, description="Log traffic at this hop")
    ]

    allowed_ports: Annotated[
        list[int],
        Field(default_factory=list, description="Allowed ports (empty = all)"),
    ]

    blocked_ports: Annotated[
        list[int],
        Field(default_factory=list, description="Blocked ports"),
    ]

    allowed_destinations: Annotated[
        list[str],
        Field(default_factory=list, description="Allowed destination IPs/domains"),
    ]

    blocked_destinations: Annotated[
        list[str],
        Field(default_factory=list, description="Blocked destination IPs/domains"),
    ]


class ChainConfig(BaseModel):
    """Configuration for creating a network chain.

    A chain defines a sequence of containers that traffic flows through,
    providing defense in depth.

    Example:
        >>> config = ChainConfig(
        ...     name="secure-browsing",
        ...     app_container="my-browser",
        ...     hops=[
        ...         ChainHop(hop_type=ChainHopType.FIREWALL, container_name="sys-firewall", position=1),
        ...         ChainHop(hop_type=ChainHopType.GATEWAY, container_name="sys-tor", position=2),
        ...     ],
        ... )
    """

    model_config = ConfigDict(frozen=True)

    name: Annotated[
        str,
        Field(
            min_length=1,
            max_length=64,
            pattern=r"^[a-zA-Z0-9][a-zA-Z0-9_.-]*$",
            description="Unique chain name",
        ),
    ]

    app_container: Annotated[str, Field(description="Source container for traffic")]

    hops: Annotated[
        list[ChainHop],
        Field(min_length=1, description="Ordered list of hops (after app)"),
    ]

    description: Annotated[str | None, Field(default=None)]

    # Chain-wide settings
    kill_switch: Annotated[
        bool,
        Field(default=True, description="Block traffic if any hop fails"),
    ]

    dns_through_chain: Annotated[
        bool,
        Field(default=True, description="Route DNS through the chain"),
    ]

    ipv6_block: Annotated[
        bool,
        Field(default=True, description="Block IPv6 to prevent leaks"),
    ]


class ChainHopState(BaseModel):
    """Runtime state of a single hop."""

    model_config = ConfigDict(frozen=False)

    container_name: str
    """Container name."""

    hop_type: ChainHopType
    """Type of hop."""

    position: int
    """Position in chain."""

    status: str = "unknown"
    """Status: unknown, running, stopped, error."""

    ip_address: str | None = None
    """IP address of this hop."""

    bytes_in: int = 0
    """Bytes received at this hop."""

    bytes_out: int = 0
    """Bytes sent from this hop."""

    last_activity: datetime | None = None
    """Last traffic activity."""

    error_message: str | None = None
    """Error message if status is 'error'."""


class ChainState(BaseModel):
    """Runtime state of a network chain."""

    model_config = ConfigDict(frozen=False)

    name: str
    """Chain name."""

    app_container: str
    """Source container."""

    status: str = "stopped"
    """Chain status: stopped, starting, running, degraded, error."""

    created_at: datetime = Field(default_factory=datetime.now)
    """When the chain was created."""

    started_at: datetime | None = None
    """When the chain was last started."""

    hops: list[ChainHopState] = Field(default_factory=list)
    """State of each hop."""

    kill_switch_active: bool = False
    """Whether kill switch is currently blocking traffic."""

    public_ip: str | None = None
    """Public IP at end of chain (exit point)."""

    description: str | None = None
    """Chain description."""

    error_message: str | None = None
    """Error message if status is 'error'."""


class ChainLogEntry(BaseModel):
    """A log entry for traffic through a chain."""

    model_config = ConfigDict(frozen=False)

    timestamp: datetime = Field(default_factory=datetime.now)
    """When this entry was logged."""

    chain_name: str
    """Chain this log belongs to."""

    hop_position: int
    """Which hop logged this."""

    hop_container: str
    """Container that logged this."""

    direction: str = "out"
    """Traffic direction: in or out."""

    protocol: str = "tcp"
    """Protocol: tcp, udp, icmp."""

    src_ip: str | None = None
    """Source IP."""

    dst_ip: str | None = None
    """Destination IP."""

    dst_port: int | None = None
    """Destination port."""

    bytes_transferred: int = 0
    """Bytes transferred."""

    action: str = "allow"
    """Action taken: allow, block, log."""


class ChainStore(BaseModel):
    """Persisted chain registry."""

    model_config = ConfigDict(frozen=False)

    version: int = 1
    chains: dict[str, ChainState] = Field(default_factory=dict)

    def add_chain(self, chain: ChainState) -> None:
        """Add or update a chain."""
        self.chains[chain.name] = chain

    def get_chain(self, name: str) -> ChainState | None:
        """Get a chain by name."""
        return self.chains.get(name)

    def remove_chain(self, name: str) -> bool:
        """Remove a chain. Returns True if removed."""
        if name in self.chains:
            del self.chains[name]
            return True
        return False

    def list_chains(self) -> list[ChainState]:
        """List all chains."""
        return list(self.chains.values())

    def get_chains_for_container(self, container_name: str) -> list[ChainState]:
        """Get all chains that involve a container."""
        result = []
        for chain in self.chains.values():
            if chain.app_container == container_name or any(h.container_name == container_name for h in chain.hops):
                result.append(chain)
        return result
