# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Generate a draw.io diagram from import_analysis.csv
Shows all modules and their imports, grouped by container.
"""

import csv
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Dict, List, Set, Tuple
from collections import defaultdict
import html

# Paths
SCRIPT_DIR = Path(__file__).parent
INPUT_FILE = SCRIPT_DIR / "import_analysis.csv"
OUTPUT_FILE = SCRIPT_DIR / "import_diagram.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": "#d4edda",      # Mint
    "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 _getShortName(moduleName: str) -> str:
    """Get short display name for module."""
    parts = moduleName.replace("gateway.", "").split(".")
    if moduleName == "gateway.app":
        return "app"
    # Return last part(s) for readability
    if len(parts) > 2:
        return ".".join(parts[-2:])
    return parts[-1] if parts else moduleName

def _parseImports() -> Tuple[Dict[str, Set[str]], List[Tuple[str, str, str]]]:
    """
    Parse import_analysis.csv and return:
    - containers: Dict mapping container name to set of modules
    - edges: List of (source, target, label) tuples
    """
    containers = defaultdict(set)
    edges = []
    seenEdges = set()
    
    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"]
            position = row["position"]
            
            # 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
            
            # Add to containers
            containers[sourceContainer].add(sourceFull)
            containers[targetContainer].add(targetFull)
            
            # Create edge key to avoid duplicates
            edgeKey = (sourceFull, targetFull)
            if edgeKey not in seenEdges:
                seenEdges.add(edgeKey)
                # Simplify label
                label = "header" if position == "header" else position.replace("function ", "fn:")
                edges.append((sourceFull, targetFull, label))
    
    return dict(containers), edges

def _generateDrawio(containers: Dict[str, Set[str]], edges: List[Tuple[str, str, str]]) -> str:
    """Generate draw.io XML content."""
    
    # Create module ID mapping
    moduleIds = {}
    idCounter = [2]  # Start at 2 (0 and 1 are reserved)
    
    def _getModuleId(moduleName: str) -> str:
        if moduleName not in moduleIds:
            moduleIds[moduleName] = f"node_{idCounter[0]}"
            idCounter[0] += 1
        return moduleIds[moduleName]
    
    # Calculate positions
    containerX = 0
    containerY = 0
    containerWidth = 300
    containerHeight = 0
    containerPadding = 50
    moduleHeight = 30
    moduleWidth = 250
    modulePadding = 10
    
    cells = []
    
    # Sort containers
    sortedContainers = sorted(containers.keys())
    
    # Create container groups and modules
    containerPositions = {}
    currentY = 0
    currentX = 0
    maxHeightInRow = 0
    containersPerRow = 3
    containerCount = 0
    
    for containerName in sortedContainers:
        modules = sorted(containers[containerName])
        
        # Calculate container size
        numModules = len(modules)
        height = numModules * (moduleHeight + modulePadding) + 60
        
        # Position container
        if containerCount > 0 and containerCount % containersPerRow == 0:
            currentY += maxHeightInRow + containerPadding
            currentX = 0
            maxHeightInRow = 0
        
        containerPositions[containerName] = (currentX, currentY)
        maxHeightInRow = max(maxHeightInRow, height)
        
        # Get color
        baseContainer = containerName.split(".")[0]
        color = CONTAINER_COLORS.get(baseContainer, "#ffffff")
        
        # Create container (swimlane)
        containerId = f"container_{containerName.replace('.', '_')}"
        cells.append(f'''      <mxCell id="{containerId}" value="{html.escape(containerName)}" style="swimlane;fontStyle=1;childLayout=stackLayout;horizontal=1;startSize=30;horizontalStack=0;resizeParent=1;resizeParentMax=0;resizeLast=0;collapsible=1;marginBottom=0;whiteSpace=wrap;html=1;fillColor={color};" vertex="1" parent="1">
        <mxGeometry x="{currentX}" y="{currentY}" width="{containerWidth}" height="{height}" as="geometry" />
      </mxCell>''')
        
        # Create module nodes inside container
        moduleY = 30
        for moduleName in modules:
            moduleId = _getModuleId(moduleName)
            shortName = _getShortName(moduleName)
            cells.append(f'''      <mxCell id="{moduleId}" value="{html.escape(shortName)}" style="text;strokeColor=none;fillColor=none;align=left;verticalAlign=middle;spacingLeft=4;spacingRight=4;overflow=hidden;points=[[0,0.5],[1,0.5]];portConstraint=eastwest;rotatable=0;whiteSpace=wrap;html=1;fontSize=10;" vertex="1" parent="{containerId}">
        <mxGeometry y="{moduleY}" width="{containerWidth}" height="{moduleHeight}" as="geometry" />
      </mxCell>''')
            moduleY += moduleHeight + modulePadding
        
        currentX += containerWidth + containerPadding
        containerCount += 1
    
    # Create edges (only between different containers to reduce clutter)
    edgeId = idCounter[0]
    for source, target, label in edges:
        sourceContainer = _getContainer(source)
        targetContainer = _getContainer(target)
        
        # Skip internal container edges for clarity
        if sourceContainer == targetContainer:
            continue
        
        sourceId = _getModuleId(source)
        targetId = _getModuleId(target)
        
        # Shorten label if too long
        displayLabel = label[:20] + "..." if len(label) > 20 else label
        
        cells.append(f'''      <mxCell id="edge_{edgeId}" value="{html.escape(displayLabel)}" style="edgeStyle=orthogonalEdgeStyle;rounded=0;orthogonalLoop=1;jettySize=auto;html=1;fontSize=8;labelBackgroundColor=#ffffff;" 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="import-diagram" name="Module 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="3000" pageHeight="2000" 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 imports...")
    containers, edges = _parseImports()
    
    print(f"Found {len(containers)} containers with {sum(len(m) for m in containers.values())} modules")
    print(f"Found {len(edges)} unique import edges")
    
    print("Generating draw.io diagram...")
    xml = _generateDrawio(containers, edges)
    
    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        f.write(xml)
    
    print(f"Diagram saved to: {OUTPUT_FILE}")
    
    # Print container summary
    print("\nContainers:")
    for name, modules in sorted(containers.items()):
        print(f"  {name}: {len(modules)} modules")

if __name__ == "__main__":
    main()
