# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Generate a simplified draw.io diagram showing container-to-container imports.
Aggregates all module imports into container-level relationships.
"""

import csv
from pathlib import Path
from typing import Dict, Tuple
from collections import defaultdict
import html
import math

# Paths
SCRIPT_DIR = Path(__file__).parent
INPUT_FILE = SCRIPT_DIR / "import_analysis.csv"
OUTPUT_FILE = SCRIPT_DIR / "import_diagram_containers.drawio"

# Container colors
CONTAINER_COLORS = {
    "app": "#dae8fc",         # Light blue
    "aichat": "#d5e8d4",      # Light green
    "auth": "#ffe6cc",        # Light orange
    "connectors": "#e1d5e7",  # Light purple
    "datamodels": "#fff2cc",  # Light yellow
    "interfaces": "#f8cecc",  # Light red
    "routes": "#d0cee2",      # Light violet
    "security": "#fad7ac",    # Peach
    "services": "#b1ddf0",    # Sky blue
    "shared": "#f0fff0",      # Honeydew
    "workflows": "#f5f5f5",   # Light gray
    "features": "#e2efda",    # Sage green
}


def _getContainer(moduleName: str) -> str:
    """Extract container name from module path."""
    if moduleName == "gateway.app":
        return "app"
    
    parts = moduleName.replace("gateway.", "").split(".")
    if len(parts) < 2:
        return "app"
    
    container = parts[1]  # modules.XXX or tests.XXX or scripts
    
    # Skip tests and scripts
    if container in ("tests", "scripts") or container.startswith("script_"):
        return None
    if parts[0] in ("tests", "scripts"):
        return None
    
    # Handle features sub-containers
    if container == "features" and len(parts) > 2:
        return f"features.{parts[2]}"
    
    return container


def _parseContainerImports() -> Tuple[Dict[str, int], Dict[Tuple[str, str], int]]:
    """
    Parse import_analysis.csv and return:
    - containerModuleCounts: Dict mapping container name to module count
    - containerEdges: Dict mapping (source_container, target_container) to import count
    """
    containerModules = defaultdict(set)
    containerEdges = defaultdict(int)
    
    with open(INPUT_FILE, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            sourceFull = row["module_name"]
            targetFull = row["imported_module_name"]
            
            # Skip external imports and relative imports
            if not targetFull.startswith("modules."):
                continue
            if targetFull.startswith("(relative)"):
                continue
            
            # Add gateway prefix to target for consistency
            targetFull = f"gateway.{targetFull}"
            
            # Get containers
            sourceContainer = _getContainer(sourceFull)
            targetContainer = _getContainer(targetFull)
            
            # Skip if either is None (tests/scripts)
            if sourceContainer is None or targetContainer is None:
                continue
            
            # Track modules per container
            containerModules[sourceContainer].add(sourceFull)
            containerModules[targetContainer].add(targetFull)
            
            # Skip self-imports (within same container)
            if sourceContainer == targetContainer:
                continue
            
            # Count container-to-container imports
            containerEdges[(sourceContainer, targetContainer)] += 1
    
    # Convert module sets to counts
    containerModuleCounts = {k: len(v) for k, v in containerModules.items()}
    
    return containerModuleCounts, dict(containerEdges)


def _generateDrawio(containerModuleCounts: Dict[str, int], 
                    containerEdges: Dict[Tuple[str, str], int]) -> str:
    """Generate draw.io XML content with container nodes and aggregated edges."""
    
    containers = sorted(containerModuleCounts.keys())
    numContainers = len(containers)
    
    # Arrange containers in a circle for better visibility
    centerX = 600
    centerY = 500
    radius = 400
    nodeWidth = 140
    nodeHeight = 60
    
    # Calculate positions
    containerPositions = {}
    for i, container in enumerate(containers):
        angle = (2 * math.pi * i / numContainers) - math.pi / 2  # Start from top
        x = centerX + radius * math.cos(angle) - nodeWidth / 2
        y = centerY + radius * math.sin(angle) - nodeHeight / 2
        containerPositions[container] = (int(x), int(y))
    
    cells = []
    
    # Create container nodes
    for container in containers:
        x, y = containerPositions[container]
        moduleCount = containerModuleCounts[container]
        
        # Get color
        baseContainer = container.split(".")[0]
        color = CONTAINER_COLORS.get(baseContainer, "#ffffff")
        
        # Create node
        label = f"{container}\\n({moduleCount} modules)"
        cells.append(f'''      <mxCell id="container_{container.replace('.', '_')}" value="{html.escape(label)}" style="rounded=1;whiteSpace=wrap;html=1;fillColor={color};strokeColor=#666666;fontStyle=1;fontSize=11;" vertex="1" parent="1">
        <mxGeometry x="{x}" y="{y}" width="{nodeWidth}" height="{nodeHeight}" as="geometry" />
      </mxCell>''')
    
    # Create edges with import counts
    edgeId = 1000
    for (source, target), count in sorted(containerEdges.items(), key=lambda x: -x[1]):
        sourceId = f"container_{source.replace('.', '_')}"
        targetId = f"container_{target.replace('.', '_')}"
        
        # Thicker line for more imports
        strokeWidth = min(1 + count // 10, 5)
        
        cells.append(f'''      <mxCell id="edge_{edgeId}" value="{count}" style="edgeStyle=orthogonalEdgeStyle;rounded=1;orthogonalLoop=1;jettySize=auto;html=1;strokeWidth={strokeWidth};fontSize=10;labelBackgroundColor=#ffffff;fontStyle=1;" edge="1" parent="1" source="{sourceId}" target="{targetId}">
        <mxGeometry relative="1" as="geometry" />
      </mxCell>''')
        edgeId += 1
    
    # Assemble XML
    xml = f'''<?xml version="1.0" encoding="UTF-8"?>
<mxfile host="app.diagrams.net" modified="2025-01-22T00:00:00.000Z" agent="Python Script" version="21.0.0" type="device">
  <diagram id="container-diagram" name="Container Imports">
    <mxGraphModel dx="1434" dy="780" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="1600" pageHeight="1200" math="0" shadow="0">
      <root>
        <mxCell id="0" />
        <mxCell id="1" parent="0" />
{chr(10).join(cells)}
      </root>
    </mxGraphModel>
  </diagram>
</mxfile>'''
    
    return xml


def main():
    """Main function."""
    print("Parsing container imports...")
    containerModuleCounts, containerEdges = _parseContainerImports()
    
    print(f"Found {len(containerModuleCounts)} containers")
    print(f"Found {len(containerEdges)} container-to-container relationships")
    
    # Print summary
    print("\nContainer Import Summary:")
    print("-" * 60)
    
    # Sort by total imports (outgoing)
    outgoingCounts = defaultdict(int)
    incomingCounts = defaultdict(int)
    for (source, target), count in containerEdges.items():
        outgoingCounts[source] += count
        incomingCounts[target] += count
    
    for container in sorted(containerModuleCounts.keys()):
        modules = containerModuleCounts[container]
        outgoing = outgoingCounts.get(container, 0)
        incoming = incomingCounts.get(container, 0)
        print(f"  {container:25} {modules:3} modules | imports: {outgoing:4} out, {incoming:4} in")
    
    print("\nTop 15 Container Dependencies:")
    print("-" * 60)
    sortedEdges = sorted(containerEdges.items(), key=lambda x: -x[1])[:15]
    for (source, target), count in sortedEdges:
        print(f"  {source:25} -> {target:25} : {count:4} imports")
    
    print("\nGenerating draw.io diagram...")
    xml = _generateDrawio(containerModuleCounts, containerEdges)
    
    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        f.write(xml)
    
    print(f"\nDiagram saved to: {OUTPUT_FILE}")


if __name__ == "__main__":
    main()
