agent/topology_diagram.py

"""
Generate network topology diagrams from GNS3 topology data
"""

from typing import Dict, List, Optional
import json


def generate_ascii_diagram(topology: Dict) -> str:
    """
    Generate a simple ASCII diagram from topology data
    
    Args:
        topology: Topology dict with nodes and links
        
    Returns:
        ASCII art string representing the network
    """
    nodes = topology.get('nodes', [])
    links = topology.get('links', [])
    
    if not nodes:
        return "No nodes in topology"
    
    # Simple representation
    lines = []
    lines.append(f"{'='*60}")
    lines.append(f"Network Topology: {topology.get('project', 'Unknown')}")
    lines.append(f"Status: {topology.get('status', 'unknown')}")
    lines.append(f"{'='*60}")
    lines.append("")
    
    # Group nodes by type
    by_type = {}
    for node in nodes:
        node_type = node.get('node_type', 'unknown')
        if node_type not in by_type:
            by_type[node_type] = []
        by_type[node_type].append(node)
    
    # Display nodes by type
    for node_type, type_nodes in by_type.items():
        lines.append(f"\n{node_type.upper()} Devices ({len(type_nodes)}):")
        lines.append("-" * 60)
        for node in type_nodes:
            status = node.get('status', 'unknown')
            console = node.get('console', 'N/A')
            status_icon = "🟢" if status == "started" else "🔴" if status == "stopped" else "⚪"
            lines.append(f"  {status_icon} {node['name']:<30} Console: {console}")
    
    # Display link count
    lines.append(f"\n\nTotal Links: {len(links)}")
    
    return "\n".join(lines)


def generate_mermaid_diagram(topology: Dict) -> str:
    """
    Generate a Mermaid diagram from topology data
    
    Args:
        topology: Topology dict with nodes and links
        
    Returns:
        Mermaid diagram code
    """
    nodes = topology.get('nodes', [])
    links = topology.get('links', [])
    
    if not nodes:
        return "graph TD\n  A[No Topology Data]"
    
    lines = []
    lines.append("graph TD")
    
    # Create node definitions with icons based on type
    node_map = {}
    for i, node in enumerate(nodes):
        node_id = f"N{i}"
        node_map[node['name']] = node_id
        node_type = node.get('node_type', 'unknown')
        
        # Choose icon based on type
        if node_type == 'qemu' or 'SW' in node['name']:
            icon = "🔀"  # Switch
        elif node_type == 'vpcs' or 'PC' in node['name'] or 'POS' in node['name']:
            icon = "💻"  # PC
        elif node_type == 'cloud':
            icon = "☁️"  # Cloud
        else:
            icon = "📦"
        
        status = node.get('status', 'unknown')
        if status == "started":
            lines.append(f'  {node_id}["{icon} {node["name"]}<br/>🟢 Running"]')
        elif status == "stopped":
            lines.append(f'  {node_id}["{icon} {node["name"]}<br/>🔴 Stopped"]')
        else:
            lines.append(f'  {node_id}["{icon} {node["name"]}"]')
    
    # Add links (simplified - just show connections exist)
    # Note: GNS3 link format is complex, this is a simplified version
    if links:
        lines.append("\n  %% Network Links")
        # For now, just note that links exist
        lines.append(f"  %% {len(links)} links configured")
    
    return "\n".join(lines)


def generate_topology_summary(topology: Dict) -> str:
    """
    Generate a text summary of the topology
    
    Args:
        topology: Topology dict with nodes and links
        
    Returns:
        Markdown-formatted summary
    """
    nodes = topology.get('nodes', [])
    links = topology.get('links', [])
    project = topology.get('project', 'Unknown')
    status = topology.get('status', 'unknown')
    
    # Count nodes by type and status
    by_type = {}
    by_status = {'started': 0, 'stopped': 0, 'unknown': 0}
    
    for node in nodes:
        node_type = node.get('node_type', 'unknown')
        by_type[node_type] = by_type.get(node_type, 0) + 1
        
        node_status = node.get('status', 'unknown')
        by_status[node_status] = by_status.get(node_status, 0) + 1
    
    summary = []
    summary.append(f"## Network Topology: {project}")
    summary.append(f"**Project Status:** {status}")
    summary.append("")
    summary.append(f"### Devices ({len(nodes)} total)")
    
    for node_type, count in sorted(by_type.items()):
        summary.append(f"- **{node_type}**: {count} device(s)")
    
    summary.append("")
    summary.append(f"### Device Status")
    summary.append(f"- 🟢 Running: {by_status.get('started', 0)}")
    summary.append(f"- 🔴 Stopped: {by_status.get('stopped', 0)}")
    summary.append(f"- ⚪ Unknown: {by_status.get('unknown', 0)}")
    
    summary.append("")
    summary.append(f"### Connectivity")
    summary.append(f"- **Links**: {len(links)} connection(s)")
    
    # List console ports for access
    summary.append("")
    summary.append("### Device Access")
    console_devices = [n for n in nodes if n.get('console')]
    if console_devices:
        summary.append("Devices with console access:")
        for node in console_devices[:10]:  # Limit to first 10
            console = node.get('console')
            console_type = node.get('console_type', 'telnet')
            summary.append(f"- **{node['name']}**: {console_type}://localhost:{console}")
    
    return "\n".join(summary)