"""Mingrammer diagrams generator for cloud infrastructure diagrams."""

from typing import Dict, List, Set, Tuple
from .base import DiagramGenerator
from ..models import DiagramData, Component, Connection


class MingrammerGenerator(DiagramGenerator):
    """Generator for mingrammer/diagrams (Python code for cloud diagrams)."""

    # Mapping of component types to diagrams library imports
    AWS_MAPPINGS = {
        "aws_vpc": ("diagrams.aws.network", "VPC"),
        "aws_subnet": ("diagrams.aws.network", "PrivateSubnet"),
        "aws_ec2": ("diagrams.aws.compute", "EC2"),
        "aws_ecs": ("diagrams.aws.compute", "ECS"),
        "aws_eks": ("diagrams.aws.compute", "EKS"),
        "aws_lambda": ("diagrams.aws.compute", "Lambda"),
        "aws_rds": ("diagrams.aws.database", "RDS"),
        "aws_dynamodb": ("diagrams.aws.database", "Dynamodb"),
        "aws_s3": ("diagrams.aws.storage", "S3"),
        "aws_elb": ("diagrams.aws.network", "ELB"),
        "aws_alb": ("diagrams.aws.network", "ELB"),
        "aws_nlb": ("diagrams.aws.network", "NLB"),
        "aws_api_gateway": ("diagrams.aws.network", "APIGateway"),
        "aws_cloudfront": ("diagrams.aws.network", "CloudFront"),
        "aws_route53": ("diagrams.aws.network", "Route53"),
        "aws_security_group": ("diagrams.aws.network", "Nacl"),
        "aws_internet_gateway": ("diagrams.aws.network", "InternetGateway"),
        "aws_nat_gateway": ("diagrams.aws.network", "NATGateway"),
        "aws_elasticache": ("diagrams.aws.database", "ElastiCache"),
        "aws_sqs": ("diagrams.aws.integration", "SQS"),
        "aws_sns": ("diagrams.aws.integration", "SNS"),
        "aws_cloudwatch": ("diagrams.aws.management", "Cloudwatch"),
    }

    AZURE_MAPPINGS = {
        "azure_vm": ("diagrams.azure.compute", "VM"),
        "azure_aks": ("diagrams.azure.compute", "AKS"),
        "azure_functions": ("diagrams.azure.compute", "FunctionApps"),
        "azure_sql": ("diagrams.azure.database", "SQLDatabases"),
        "azure_cosmos": ("diagrams.azure.database", "CosmosDb"),
        "azure_storage": ("diagrams.azure.storage", "BlobStorage"),
        "azure_load_balancer": ("diagrams.azure.network", "LoadBalancers"),
        "azure_app_gateway": ("diagrams.azure.network", "ApplicationGateway"),
        "azure_vnet": ("diagrams.azure.network", "VirtualNetworks"),
    }

    GCP_MAPPINGS = {
        "gcp_compute": ("diagrams.gcp.compute", "ComputeEngine"),
        "gcp_gke": ("diagrams.gcp.compute", "GKE"),
        "gcp_functions": ("diagrams.gcp.compute", "Functions"),
        "gcp_cloud_sql": ("diagrams.gcp.database", "SQL"),
        "gcp_firestore": ("diagrams.gcp.database", "Firestore"),
        "gcp_storage": ("diagrams.gcp.storage", "GCS"),
        "gcp_load_balancer": ("diagrams.gcp.network", "LoadBalancing"),
        "gcp_vpc": ("diagrams.gcp.network", "VPC"),
    }

    K8S_MAPPINGS = {
        "k8s_pod": ("diagrams.k8s.compute", "Pod"),
        "k8s_deployment": ("diagrams.k8s.compute", "Deploy"),
        "k8s_service": ("diagrams.k8s.network", "Service"),
        "k8s_ingress": ("diagrams.k8s.network", "Ingress"),
        "k8s_configmap": ("diagrams.k8s.config", "ConfigMap"),
        "k8s_secret": ("diagrams.k8s.config", "Secret"),
        "k8s_pv": ("diagrams.k8s.storage", "PV"),
        "k8s_pvc": ("diagrams.k8s.storage", "PVC"),
    }

    GENERIC_MAPPINGS = {
        "compute": ("diagrams.generic.compute", "Rack"),
        "database": ("diagrams.generic.database", "SQL"),
        "storage": ("diagrams.generic.storage", "Storage"),
        "network": ("diagrams.generic.network", "Router"),
        "load_balancer": ("diagrams.generic.network", "LoadBalancer"),
    }

    def __init__(self, diagram_data: DiagramData = None):
        """Initialize Mingrammer generator.

        Args:
            diagram_data: DiagramData model
        """
        super().__init__(diagram_data)
        self.imports: Set[Tuple[str, str]] = set()
        self.clusters: Dict[str, List[Component]] = {}

    def get_format(self) -> str:
        """Return format name."""
        return "python-diagrams"

    def get_file_extension(self) -> str:
        """Return file extension."""
        return ".py"

    def generate(self) -> str:
        """Generate Python code using mingrammer/diagrams.

        Returns:
            Python code string
        """
        # Reset state
        self.imports = set()
        self.clusters = {}

        # Organize components into clusters
        self._organize_clusters()

        # Build the Python code
        code_parts = []

        # Add header
        code_parts.append('"""Generated by SnapInfra - Architecture Diagram"""')
        code_parts.append("")

        # Add imports
        code_parts.append(self._generate_imports())
        code_parts.append("")

        # Add main diagram code
        code_parts.append(self._generate_diagram_code())

        return "\n".join(code_parts)

    def _organize_clusters(self) -> None:
        """Organize components into logical clusters."""
        # Group components by cluster property or by type
        for component in self.diagram_data.components:
            cluster_name = component.properties.get("cluster", "main")

            if cluster_name not in self.clusters:
                self.clusters[cluster_name] = []

            self.clusters[cluster_name].append(component)

    def _get_import_for_component(self, component: Component) -> Tuple[str, str]:
        """Get the import module and class for a component.

        Args:
            component: Component to get import for

        Returns:
            Tuple of (module, class_name)
        """
        type_lower = component.type.lower()

        # Try AWS mappings
        if type_lower in self.AWS_MAPPINGS:
            return self.AWS_MAPPINGS[type_lower]

        # Try Azure mappings
        if type_lower in self.AZURE_MAPPINGS:
            return self.AZURE_MAPPINGS[type_lower]

        # Try GCP mappings
        if type_lower in self.GCP_MAPPINGS:
            return self.GCP_MAPPINGS[type_lower]

        # Try K8s mappings
        if type_lower in self.K8S_MAPPINGS:
            return self.K8S_MAPPINGS[type_lower]

        # Try generic mappings
        if type_lower in self.GENERIC_MAPPINGS:
            return self.GENERIC_MAPPINGS[type_lower]

        # Default to generic compute
        return ("diagrams.generic.compute", "Rack")

    def _generate_imports(self) -> str:
        """Generate import statements.

        Returns:
            Import statements as string
        """
        # Always import base Diagram and Cluster
        imports = ["from diagrams import Diagram, Cluster, Edge"]

        # Collect all needed imports from components
        import_map: Dict[str, Set[str]] = {}

        for component in self.diagram_data.components:
            module, class_name = self._get_import_for_component(component)
            self.imports.add((module, class_name))

            if module not in import_map:
                import_map[module] = set()
            import_map[module].add(class_name)

        # Generate import statements
        for module in sorted(import_map.keys()):
            classes = ", ".join(sorted(import_map[module]))
            imports.append(f"from {module} import {classes}")

        return "\n".join(imports)

    def _generate_diagram_code(self) -> str:
        """Generate the main diagram code.

        Returns:
            Diagram code as string
        """
        code_parts = []

        # Get diagram name from metadata or use default
        diagram_name = self.diagram_data.metadata.get("name", "Infrastructure Architecture")
        diagram_direction = self.diagram_data.metadata.get("direction", "LR")

        code_parts.append(f'with Diagram("{diagram_name}", show=False, direction="{diagram_direction}"):')

        # Generate components and clusters
        if len(self.clusters) > 1 or (len(self.clusters) == 1 and "main" not in self.clusters):
            # Multiple clusters or non-main cluster
            code_parts.extend(self._generate_clustered_components())
        else:
            # Single main cluster or no clusters
            code_parts.extend(self._generate_flat_components())

        # Generate connections
        code_parts.append("")
        code_parts.append("    # Connections")
        code_parts.extend(self._generate_connections())

        return "\n".join(code_parts)

    def _generate_flat_components(self) -> List[str]:
        """Generate components without clusters.

        Returns:
            List of code lines
        """
        lines = []

        for component in self.diagram_data.components:
            var_name = self._sanitize_identifier(component.name)
            _, class_name = self._get_import_for_component(component)
            label = component.name

            lines.append(f'    {var_name} = {class_name}("{label}")')

        return lines

    def _generate_clustered_components(self) -> List[str]:
        """Generate components organized in clusters.

        Returns:
            List of code lines
        """
        lines = []

        for cluster_name, components in self.clusters.items():
            # Create cluster
            lines.append(f'    with Cluster("{cluster_name}"):')

            # Add components to cluster
            for component in components:
                var_name = self._sanitize_identifier(component.name)
                _, class_name = self._get_import_for_component(component)
                label = component.name

                lines.append(f'        {var_name} = {class_name}("{label}")')

            lines.append("")  # Empty line after cluster

        return lines

    def _generate_connections(self) -> List[str]:
        """Generate connection statements.

        Returns:
            List of code lines
        """
        lines = []

        for connection in self.diagram_data.connections:
            source_component = self.get_component_by_id(connection.source)
            target_component = self.get_component_by_id(connection.target)

            if not source_component or not target_component:
                continue

            source_var = self._sanitize_identifier(source_component.name)
            target_var = self._sanitize_identifier(target_component.name)

            # Check if connection has a label
            label = connection.properties.get("label", "")

            if label:
                lines.append(f'    {source_var} >> Edge(label="{label}") >> {target_var}')
            else:
                lines.append(f'    {source_var} >> {target_var}')

        return lines

    def generate_example(self) -> str:
        """Generate an example 3-tier AWS architecture.

        Returns:
            Python code string
        """
        return '''"""Generated by SnapInfra - 3-Tier AWS Architecture Example"""

from diagrams import Diagram, Cluster, Edge
from diagrams.aws.compute import ECS
from diagrams.aws.database import RDS
from diagrams.aws.network import ELB, VPC
from diagrams.aws.storage import S3

with Diagram("3-Tier AWS Architecture", show=False, direction="LR"):
    with Cluster("VPC (10.0.0.0/16)"):
        with Cluster("Public Subnet"):
            alb = ELB("Application Load Balancer")

        with Cluster("Private Subnet - App Tier"):
            app_servers = [
                ECS("App Server 1"),
                ECS("App Server 2"),
                ECS("App Server 3")
            ]

        with Cluster("Private Subnet - Data Tier"):
            primary_db = RDS("Primary DB")
            replica_db = RDS("Read Replica")

    storage = S3("Static Assets")

    # Connections
    alb >> Edge(label="http") >> app_servers
    for app in app_servers:
        app >> Edge(label="sql") >> primary_db
        app >> Edge(label="s3 api") >> storage

    primary_db >> Edge(label="replication") >> replica_db
'''
