"""Data models for infra-diagram."""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum


class NodeType(str, Enum):
    """Infrastructure component types."""
    # Compute
    COMPUTE = "compute"
    CONTAINER = "container"
    SERVERLESS = "serverless"
    VM = "vm"

    # Storage
    STORAGE = "storage"
    DATABASE = "database"
    CACHE = "cache"
    QUEUE = "queue"

    # Network
    NETWORK = "network"
    LOAD_BALANCER = "load_balancer"
    CDN = "cdn"
    DNS = "dns"
    FIREWALL = "firewall"
    VPN = "vpn"
    GATEWAY = "gateway"
    SUBNET = "subnet"

    # Security
    IAM = "iam"
    SECRET = "secret"
    CERTIFICATE = "certificate"

    # Monitoring
    MONITORING = "monitoring"
    LOGGING = "logging"

    # Other
    SERVICE = "service"
    VOLUME = "volume"
    CONFIG = "config"
    REGISTRY = "registry"
    INGRESS = "ingress"
    NAMESPACE = "namespace"
    DEPLOYMENT = "deployment"
    POD = "pod"
    UNKNOWN = "unknown"


class SourceType(str, Enum):
    """Source file type."""
    TERRAFORM = "terraform"
    KUBERNETES = "kubernetes"
    DOCKER_COMPOSE = "docker-compose"
    UNKNOWN = "unknown"


# Icons for each node type (used in Mermaid and HTML)
NODE_ICONS: dict[NodeType, str] = {
    NodeType.COMPUTE: "🖥️",
    NodeType.CONTAINER: "🐳",
    NodeType.SERVERLESS: "⚡",
    NodeType.VM: "💻",
    NodeType.STORAGE: "💾",
    NodeType.DATABASE: "🗄️",
    NodeType.CACHE: "⚡",
    NodeType.QUEUE: "📬",
    NodeType.NETWORK: "🌐",
    NodeType.LOAD_BALANCER: "⚖️",
    NodeType.CDN: "🌍",
    NodeType.DNS: "📡",
    NodeType.FIREWALL: "🛡️",
    NodeType.VPN: "🔒",
    NodeType.GATEWAY: "🚪",
    NodeType.SUBNET: "🔗",
    NodeType.IAM: "👤",
    NodeType.SECRET: "🔑",
    NodeType.CERTIFICATE: "📜",
    NodeType.MONITORING: "📊",
    NodeType.LOGGING: "📝",
    NodeType.SERVICE: "⚙️",
    NodeType.VOLUME: "📦",
    NodeType.CONFIG: "⚙️",
    NodeType.REGISTRY: "📦",
    NodeType.INGRESS: "🚪",
    NodeType.NAMESPACE: "📁",
    NodeType.DEPLOYMENT: "🚀",
    NodeType.POD: "🐳",
    NodeType.UNKNOWN: "❓",
}


@dataclass
class InfraNode:
    """A single infrastructure component."""
    id: str
    name: str
    node_type: NodeType
    provider: str = ""  # aws, azure, gcp, k8s, docker
    source_file: str = ""
    properties: dict = field(default_factory=dict)
    children: list[str] = field(default_factory=list)  # child node IDs

    @property
    def icon(self) -> str:
        return NODE_ICONS.get(self.node_type, "❓")

    @property
    def label(self) -> str:
        return f"{self.icon} {self.name}"

    @property
    def safe_id(self) -> str:
        """Mermaid-safe ID."""
        return self.id.replace("-", "_").replace(".", "_").replace("/", "_")


@dataclass
class InfraEdge:
    """A connection between two infrastructure components."""
    source: str  # node ID
    target: str  # node ID
    label: str = ""
    style: str = ""  # solid, dashed, dotted


@dataclass
class InfraGroup:
    """A group/subgraph of related components."""
    id: str
    name: str
    node_ids: list[str] = field(default_factory=list)
    style: str = ""


@dataclass
class InfraDiagram:
    """Complete infrastructure diagram."""
    title: str = "Infrastructure Diagram"
    source_type: SourceType = SourceType.UNKNOWN
    nodes: list[InfraNode] = field(default_factory=list)
    edges: list[InfraEdge] = field(default_factory=list)
    groups: list[InfraGroup] = field(default_factory=list)
    source_files: list[str] = field(default_factory=list)

    @property
    def node_count(self) -> int:
        return len(self.nodes)

    @property
    def edge_count(self) -> int:
        return len(self.edges)

    def get_node(self, node_id: str) -> InfraNode | None:
        for n in self.nodes:
            if n.id == node_id:
                return n
        return None

    def add_node(self, node: InfraNode) -> None:
        if not self.get_node(node.id):
            self.nodes.append(node)

    def add_edge(self, edge: InfraEdge) -> None:
        self.edges.append(edge)

    def add_group(self, group: InfraGroup) -> None:
        self.groups.append(group)
