#!/usr/bin/env python3
"""
RTLViz MCP Server

Provides:
- Resource: rtlviz://prompt - The RTL diagram generation prompt
- Resource: rtlviz://skill - The skill documentation
- Tool: render_diagram - Converts DOT code to an interactive HTML viewer
- Tool: generate_rtl_diagram - Generate diagram from Verilog source

Architecture: Uses MCP Sampling for LLM - no API key required.
"""

import os
import webbrowser
from pathlib import Path
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Resource, Tool, TextContent

# Import telemetry
from rtlviz.telemetry import (
    ping_session_start,
    ping_resource_read,
    ping_tool_call,
    ping_diagram_rendered
)

# Import RTL processing modules
from rtlviz.rtl_parser import parse_project, get_connections, Module
from rtlviz.dot_generator import generate_dot
from rtlviz.llm_analyzer import enhance_instances
from rtlviz.diagram_validator import validate_diagram

# Get the assets directory (relative to this file)
ASSETS_DIR = Path(__file__).parent / "assets"

# Create the MCP server
server = Server("rtlviz")

# Send telemetry ping on import (server start)
ping_session_start()


@server.list_resources()
async def list_resources() -> list[Resource]:
    """List available resources."""
    return [
        Resource(
            uri="rtlviz://guidelines",
            name="RTL Diagram Guidelines",
            description="Comprehensive guidelines for generating high-quality RTL diagrams. READ THIS before calling render_diagram.",
            mimeType="text/markdown",
        ),
    ]


@server.read_resource()
async def read_resource(uri: str) -> str:
    """Read a resource by URI."""
    # Convert URI to string in case MCP passes AnyUrl object
    uri_str = str(uri)
    
    if uri_str == "rtlviz://guidelines":
        ping_resource_read("guidelines")
        guidelines_path = Path(__file__).parent / "PROMPT.md"
        if guidelines_path.exists():
            return guidelines_path.read_text(encoding="utf-8")
        return "Error: PROMPT.md (guidelines) not found"
    else:
        return f"Unknown resource: {uri_str}"


@server.list_tools()
async def list_tools() -> list[Tool]:
    """List available tools."""
    return [
        Tool(
            name="render_diagram",
            description="""Convert Graphviz DOT code to an interactive HTML viewer with pan/zoom.

**YOU MUST PROVIDE DOT CODE. This tool does NOT auto-generate diagrams.**

**REQUIRED WORKFLOW:**
1. READ the `rtlviz://guidelines` resource FIRST - it contains critical rules
2. Read all .v files in the source directory
3. Parse: modules, instances, registers, clock domains, CDC paths
4. Generate detailed DOT with white-box internals (registers, logic, FFs)
5. Call this tool with your DOT code

**CRITICAL RULES (from guidelines):**
- EVERY node MUST be connected - no orphan nodes
- Use `taillabel="signal  "` or `headlabel="  signal"` - NEVER xlabel
- Show internal detail: `{ reg_name | [N:0] }`, `{ +1 Logic | !full }`
- CDC sync FFs: `{ q1 | FF1 }` -> `{ q2 | FF2 }`
- Must include `splines=ortho` for orthogonal routing

**DOT TEMPLATE:**
```
digraph "Name" {
    graph [rankdir=LR, splines=ortho, nodesep=1.0, ranksep=2.0, sep="+25"];
    node [shape=record, style="filled,rounded", fontsize=10];
    // Clusters for clock domains, modules, I/O
    // Registers: fillcolor="#FFF9C4"
    // Logic: fillcolor="#E8F5E9"
    // CDC FFs: fillcolor="#FFE0B2"
}
```

**VALIDATION:** This tool validates your DOT code and will reject:
- Empty or too-short DOT code
- Missing required graph attributes
- Diagrams with too few connections""",
            inputSchema={
                "type": "object",
                "properties": {
                    "dot_content": {
                        "type": "string",
                        "description": "The Graphviz DOT code to render. REQUIRED - read guidelines first and generate proper DOT.",
                    },
                    "output_path": {
                        "type": "string",
                        "description": "Absolute path for the output HTML file",
                    },
                    "title": {
                        "type": "string",
                        "description": "Title for the diagram (default: RTL Block Diagram)",
                        "default": "RTL Block Diagram",
                    },
                },
                "required": ["dot_content", "output_path"],
            },
        ),
        Tool(
            name="generate_rtl_diagram",
            description="""Generate a DETAILED RTL block diagram from Verilog source using LLM analysis.

**This tool automatically:**
1. Reads all .v files in the source directory
2. Uses LLM (via MCP Sampling) to analyze the RTL and generate detailed DOT
3. Renders an interactive HTML with hover, zoom, and pan features

**For manual control:** Use `render_diagram` with your own DOT code instead.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "source_dir": {
                        "type": "string",
                        "description": "Absolute path to the directory containing Verilog (.v) source files",
                    },
                    "output_path": {
                        "type": "string",
                        "description": "Absolute path for the output HTML file",
                    },
                    "title": {
                        "type": "string",
                        "description": "Title for the diagram",
                        "default": "RTL Block Diagram",
                    },
                    "top_module": {
                        "type": "string",
                        "description": "Optional: manually specify the top module name (auto-detected if not provided)",
                    },
                    "use_llm": {
                        "type": "boolean",
                        "description": "Use LLM for enhanced semantic analysis (via MCP Sampling)",
                        "default": True,
                    },
                },
                "required": ["source_dir", "output_path"],
            },
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    """Handle tool calls."""
    # Log FIRST before executing tool
    ping_tool_call(name)
    
    if name == "render_diagram":
        return await render_diagram(
            dot_content=arguments.get("dot_content", ""),
            output_path=arguments["output_path"],
            title=arguments.get("title", "RTL Block Diagram"),
        )
    elif name == "generate_rtl_diagram":
        return await generate_rtl_diagram(
            source_dir=arguments["source_dir"],
            output_path=arguments["output_path"],
            title=arguments.get("title", "RTL Block Diagram"),
            top_module=arguments.get("top_module"),
            use_llm=arguments.get("use_llm", True),
        )
    else:
        return [TextContent(type="text", text=f"Unknown tool: {name}")]


def validate_agent_dot(dot_code: str) -> tuple[bool, list[str]]:
    """Validate DOT code meets minimum quality standards for agent-generated diagrams.
    
    Returns:
        tuple: (is_valid, list_of_errors)
    """
    errors = []
    
    # Check for empty or very short DOT
    if not dot_code or len(dot_code.strip()) < 100:
        errors.append("DOT code is empty or too short (minimum 100 characters).")
        errors.append("You must read rtlviz://guidelines and generate proper DOT code.")
        return False, errors
    
    # Check for digraph declaration
    if "digraph" not in dot_code.lower():
        errors.append("Missing 'digraph' declaration. DOT must start with 'digraph \"Name\" {'")
    
    # Check for required graph attributes
    if "splines=ortho" not in dot_code and "splines = ortho" not in dot_code:
        errors.append("Missing 'splines=ortho'. Required for orthogonal line routing.")
    
    # Check for minimum connections (arrows)
    arrow_count = dot_code.count("->")
    if arrow_count < 3:
        errors.append(f"Too few connections ({arrow_count} found, minimum 3). All nodes must be connected.")
    
    # Check for node definitions
    if "[" not in dot_code or "]" not in dot_code:
        errors.append("No node attributes found. Nodes should have labels and colors.")
    
    if errors:
        return False, errors
    return True, []


async def render_diagram(dot_content: str, output_path: str, title: str) -> list[TextContent]:
    """Generate an interactive HTML viewer from DOT code.
    
    Validates DOT code quality before rendering to ensure agents follow guidelines.
    """
    try:
        # Validate DOT code quality
        is_valid, errors = validate_agent_dot(dot_content)
        if not is_valid:
            error_msg = "\n".join(f"  • {e}" for e in errors)
            return [
                TextContent(
                    type="text",
                    text=f"❌ **Invalid DOT code. Please follow the workflow:**\n\n"
                         f"{error_msg}\n\n"
                         f"**Required workflow:**\n"
                         f"1. Read the `rtlviz://guidelines` resource\n"
                         f"2. Read the Verilog source files (.v)\n"
                         f"3. Generate detailed DOT code following the guidelines\n"
                         f"4. Call this tool with proper DOT code\n\n"
                         f"Use `list_resources` to see available resources."
                )
            ]
        
        # Escape DOT content for JavaScript template literal
        dot_escaped = (
            dot_content
            .replace("\\", "\\\\")
            .replace("`", "\\`")
            .replace("${", "\\${")
        )

        # Generate HTML
        html = generate_html(dot_escaped, title)

        # Ensure output directory exists
        out_path = Path(output_path)
        out_path.parent.mkdir(parents=True, exist_ok=True)

        # Write the file
        out_path.write_text(html, encoding="utf-8")
        
        # Also write DOT file for reference
        dot_path = out_path.with_suffix('.dot')
        dot_path.write_text(dot_content, encoding="utf-8")

        # Telemetry: diagram rendered successfully
        ping_diagram_rendered(success=True)
        
        # Auto-open in browser
        webbrowser.open(f'file:///{out_path.resolve()}')

        return [
            TextContent(
                type="text",
                text=f"✅ **Diagram Generated Successfully**\n\n"
                     f"- HTML: `{output_path}`\n"
                     f"- DOT: `{dot_path}`\n\n"
                     f"Opened in browser automatically.",
            )
        ]
    except Exception as e:
        return [TextContent(type="text", text=f"❌ Error generating diagram: {e}")]


from rtlviz.direct_llm_generator import generate_full_dot

async def generate_rtl_diagram(
    source_dir: str, 
    output_path: str, 
    title: str,
    top_module: str = None,
    use_llm: bool = True
) -> list[TextContent]:
    """Generate an RTL block diagram from Verilog source files."""
    try:
        # If use_llm is requested, try that FIRST (Full Replacement)
        dot_code = None
        llm_attempted = False
        llm_succeeded = False
        
        if use_llm:
            llm_attempted = True
            # 1. Gather all source code
            all_source_code = ""
            v_files = list(Path(source_dir).glob("*.v"))
            if not v_files:
                 return [TextContent(type="text", text=f"❌ No .v files found in {source_dir}")]
            
            for v_file in v_files:
                all_source_code += f"\n// File: {v_file.name}\n"
                all_source_code += v_file.read_text(encoding="utf-8", errors="replace")
            
            # 2. Call Direct LLM Generator
            dot_code = await generate_full_dot(server, all_source_code, len(all_source_code.splitlines()))
            
            if dot_code:
                llm_succeeded = True
            # If LLM failed, we'll fall through to pattern-based below

        # Pattern-based fallback (or primary if use_llm was False or LLM failed)
        if not dot_code:
            # Traditional AST-based Path
            # Parse the Verilog project
            modules, detected_top = parse_project(source_dir)
            
            # Use provided top_module or fall back to auto-detected
            final_top = top_module if top_module else detected_top
            
            if not modules or not final_top or final_top not in modules:
                 return [TextContent(type="text", text="❌ Error parsing project or finding top module.")]
            
            # Generate DOT code
            dot_code = generate_dot(modules, final_top, title)
        
        # Strict Quality Validation
        validation = validate_diagram(dot_code)
        score = validation.get("score", 0)
        validation_pass = validation.get("pass", True)
        
        validation_msg = f"\n🔍 **Quality Score:** {score}/100"
        if not validation_pass:
            issues = validation.get("issues", [])
            validation_msg += "\n⚠️ **Issues Detected:**"
            for issue in issues[:3]:
                 validation_msg += f"\n- {issue}"
        else:
            validation_msg += " ✅ Passed Strict Validation"
        
        # Escape DOT content for JavaScript
        dot_escaped = (
            dot_code
            .replace("\\", "\\\\")
            .replace("`", "\\`")
            .replace("${", "\\${")
        )
        
        # Generate HTML
        html = generate_html(dot_escaped, title)
        
        # Write HTML file
        out_path = Path(output_path)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(html, encoding="utf-8")
        
        # Also write DOT file
        dot_path = out_path.with_suffix('.dot')
        dot_path.write_text(dot_code, encoding="utf-8")
        
        # Telemetry
        ping_diagram_rendered(success=True)
        
        # Auto-open in browser
        webbrowser.open(f'file:///{out_path.resolve()}')
        
        # Build response with stats
        if llm_succeeded:
            llm_status = "✅ Full LLM Generation"
        elif llm_attempted:
            llm_status = "⚡ Pattern-based (LLM fallback)"
        else:
            llm_status = "⚡ Pattern-based"
        
        return [
            TextContent(
                type="text",
                text=f"✅ **Diagram Generated ({llm_status})**\n\nOpen: `{output_path}`{validation_msg}",
            )
        ]
    except Exception as e:
        return [TextContent(type="text", text=f"❌ Error generating RTL diagram: {e}")]


def generate_html(dot_escaped: str, title: str) -> str:
    """Generate the HTML content with Viz.js for client-side rendering."""
    safe_title = title.replace('"', "&quot;")
    download_name = title.replace(" ", "_").lower() + "_diagram.svg"

    # Multiple CDN sources for reliability (primary + fallbacks)
    return f'''<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>{safe_title}</title>
    <style>
        * {{ margin: 0; padding: 0; box-sizing: border-box; }}
        body {{ font-family: 'Segoe UI', Arial, sans-serif; background: linear-gradient(135deg, #1a1a2e, #16213e); min-height: 100vh; padding: 20px; }}
        .container {{ background: white; border-radius: 16px; box-shadow: 0 20px 60px rgba(0,0,0,0.3); padding: 30px; max-width: 100%; }}
        h1 {{ color: #333; margin-bottom: 10px; font-weight: 600; }}
        .subtitle {{ color: #666; margin-bottom: 20px; font-size: 14px; }}
        .controls {{ margin-bottom: 20px; display: flex; gap: 10px; flex-wrap: wrap; align-items: center; }}
        button {{ padding: 10px 20px; border: none; border-radius: 8px; background: linear-gradient(135deg, #667eea, #764ba2); color: white; cursor: pointer; font-weight: 500; transition: all 0.2s; }}
        button:hover {{ transform: translateY(-2px); box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4); }}
        .zoom-info {{ font-size: 12px; color: #666; margin-left: 10px; }}
        
        /* Diagram container with scrollbars */
        #diagram {{ 
            border: 1px solid #eee; border-radius: 8px 8px 0 0; padding: 20px; 
            background: #fafafa; overflow: auto; min-height: 500px; max-height: 70vh;
            position: relative; cursor: crosshair; border-bottom: none;
        }}
        #diagram svg {{ transform-origin: top left; }}
        .loading {{ text-align: center; padding: 40px; color: #666; }}
        .error {{ color: #d32f2f; padding: 20px; background: #ffebee; border-radius: 8px; white-space: pre-wrap; }}
        .error a {{ color: #1565C0; }}
        
        /* Selection box for drag-to-zoom */
        .selection-box {{
            position: absolute; border: 2px dashed #667eea; background: rgba(102, 126, 234, 0.1);
            pointer-events: none; z-index: 100;
        }}
        
        /* Hover highlighting styles */
        #diagram svg .node:hover {{ cursor: pointer; }}
        #diagram svg .node:hover polygon, #diagram svg .node:hover ellipse, #diagram svg .node:hover path {{
            filter: drop-shadow(0 0 8px #667eea) drop-shadow(0 0 12px #764ba2);
            stroke: #667eea !important;
            stroke-width: 2.5px !important;
        }}
        #diagram.highlighting .node {{ opacity: 0.3; transition: opacity 0.2s; }}
        #diagram.highlighting .edge {{ opacity: 0.15; transition: opacity 0.2s; }}
        #diagram.highlighting .node.highlighted {{ opacity: 1; }}
        #diagram.highlighting .edge.highlighted {{ opacity: 1; stroke-width: 2px; }}
        
        /* Tooltip */
        .node-tooltip {{
            position: absolute; background: #333; color: white; padding: 8px 12px;
            border-radius: 6px; font-size: 12px; pointer-events: none; z-index: 1000;
            box-shadow: 0 4px 12px rgba(0,0,0,0.3); max-width: 300px;
            transform: translate(-50%, -100%); margin-top: -10px;
        }}
        .node-tooltip::after {{
            content: ''; position: absolute; top: 100%; left: 50%;
            border: 6px solid transparent; border-top-color: #333; transform: translateX(-50%);
        }}
        
        /* Footer Panel - continuous with diagram */
        .footer-panel {{ margin-top: 0; background: #f8f9fa; border-radius: 0 0 12px 12px; padding: 20px 25px; border: 1px solid #eee; border-top: none; }}
        .footer-section {{ margin-bottom: 20px; }}
        .footer-section:last-child {{ margin-bottom: 0; }}
        .footer-title {{ display: flex; align-items: center; gap: 8px; font-size: 14px; font-weight: 600; margin-bottom: 12px; color: #333; }}
        .footer-title::before {{ content: ''; width: 4px; height: 16px; background: #667eea; border-radius: 2px; }}
        
        /* Diagram Legend */
        .legend-grid {{ display: flex; flex-wrap: wrap; gap: 12px; }}
        .legend-item {{ display: flex; align-items: center; gap: 8px; font-size: 12px; background: white; padding: 6px 12px; border-radius: 6px; border: 1px solid #eee; color: #333; }}
        .legend-color {{ width: 20px; height: 14px; border-radius: 3px; border: 1px solid #ddd; }}
        
        /* Signal Types */
        .signal-grid {{ display: flex; flex-wrap: wrap; gap: 20px; }}
        .signal-item {{ display: flex; align-items: center; gap: 10px; font-size: 12px; color: #333; }}
        .signal-line {{ width: 40px; height: 3px; }}
        .signal-line.solid {{ background: #333; }}
        .signal-line.dashed {{ background: repeating-linear-gradient(90deg, #1565C0 0, #1565C0 6px, transparent 6px, transparent 10px); }}
        .signal-line.dotted {{ background: repeating-linear-gradient(90deg, #FF9800 0, #FF9800 3px, transparent 3px, transparent 6px); }}
        .signal-line.cdc {{ background: repeating-linear-gradient(90deg, #E91E63 0, #E91E63 8px, transparent 8px, transparent 12px); height: 4px; }}
    </style>
</head>
<body>
    <div class="container">
        <h1>{safe_title}</h1>
        <p class="subtitle">Generated with RTLViz | Hover over nodes to highlight connections</p>
        
        <div class="controls">
            <button onclick="zoomIn()">🔍 Zoom In</button>
            <button onclick="zoomOut()">🔎 Zoom Out</button>
            <button onclick="resetZoom()">↺ Reset</button>
            <button onclick="downloadSVG()">💾 Download SVG</button>
            <span class="zoom-info">Zoom: <span id="zoomLevel">100</span>% | Drag to select area for zoom</span>
        </div>
        
        <div id="diagram">
            <div class="loading">Loading Viz.js rendering engine...</div>
        </div>
        
        <!-- Footer Panel -->
        <div class="footer-panel">
            <div class="footer-section">
                <div class="footer-title">Diagram Legend</div>
                <div class="legend-grid">
                    <div class="legend-item"><span class="legend-color" style="background: #E3F2FD; border-color: #1565C0;"></span>Read Clock Domain (rclk)</div>
                    <div class="legend-item"><span class="legend-color" style="background: #FFF3E0; border-color: #E65100;"></span>Write Clock Domain (wclk)</div>
                    <div class="legend-item"><span class="legend-color" style="background: #ECEFF1;"></span>Logic Module</div>
                    <div class="legend-item"><span class="legend-color" style="background: #C8E6C9; border-color: #2E7D32;"></span>Memory (Dual-Port RAM)</div>
                    <div class="legend-item"><span class="legend-color" style="background: #FFE0B2;"></span>2-FF Synchronizers</div>
                    <div class="legend-item"><span class="legend-color" style="background: #FFF9C4;"></span>Flip-Flop Registers</div>
                </div>
            </div>
            
            <div class="footer-section">
                <div class="footer-title">Signal Types</div>
                <div class="signal-grid">
                    <div class="signal-item"><span class="signal-line solid"></span>Data Signals (wdata, rdata, addresses)</div>
                    <div class="signal-item"><span class="signal-line dotted"></span>Pointer Signals (Gray code pointers)</div>
                    <div class="signal-item"><span class="signal-line dashed"></span>Control Signals (inc, full, empty)</div>
                    <div class="signal-item"><span class="signal-line cdc"></span>Cross-Domain Signals (CDC crossing)</div>
                </div>
            </div>
        </div>
    </div>
    
    <script>
        const dotSource = `{dot_escaped}`;
        let scale = 1;
        let svgContent = '';
        
        // Multiple CDN sources for reliability
        const VIZ_SOURCES = [
            'https://unpkg.com/@viz-js/viz@3.2.4/lib/viz-standalone.js',
            'https://cdn.jsdelivr.net/npm/@viz-js/viz@3.2.4/lib/viz-standalone.js',
            'https://cdnjs.cloudflare.com/ajax/libs/viz.js/2.1.2/viz.js'
        ];
        
        async function loadVizJs() {{
            for (const src of VIZ_SOURCES) {{
                try {{
                    await new Promise((resolve, reject) => {{
                        const script = document.createElement('script');
                        script.src = src;
                        script.onload = resolve;
                        script.onerror = reject;
                        document.head.appendChild(script);
                    }});
                    console.log('Loaded Viz.js from:', src);
                    return true;
                }} catch (e) {{
                    console.warn('Failed to load from:', src);
                }}
            }}
            return false;
        }}
        
        async function renderDiagram() {{
            const diagramEl = document.getElementById('diagram');
            
            // Load Viz.js with fallbacks
            const loaded = await loadVizJs();
            if (!loaded) {{
                diagramEl.innerHTML = `<div class="error">
Failed to load Viz.js rendering engine.

This can happen due to:
1. No internet connection
2. Firewall/proxy blocking CDN
3. Browser security restrictions

Solutions:
- Check your internet connection
- Try a different browser (Chrome/Firefox recommended)
- If using file:// URL, try serving via local HTTP server

The DOT source is embedded - you can copy it and render at:
<a href="https://dreampuf.github.io/GraphvizOnline/" target="_blank">GraphvizOnline</a>
</div>`;
                return;
            }}
            
            try {{
                // Wait a bit for Viz to initialize
                await new Promise(r => setTimeout(r, 100));
                
                if (typeof Viz === 'undefined') {{
                    throw new Error('Viz not defined after load');
                }}
                
                const viz = await Viz.instance();
                const svg = viz.renderSVGElement(dotSource);
                svgContent = svg.outerHTML;
                diagramEl.innerHTML = svgContent;
                
                // Setup hover highlighting after rendering
                setupHoverHighlighting();
                setupDragToZoom();
            }} catch (error) {{
                diagramEl.innerHTML = `<div class="error">
Rendering error: ${{error.message}}

This may be due to:
1. Invalid DOT syntax
2. Browser memory limits (large diagram)

Try viewing at: <a href="https://dreampuf.github.io/GraphvizOnline/" target="_blank">GraphvizOnline</a>
</div>`;
                console.error('Viz.js render error:', error);
            }}
        }}
        
        // Build edge connection map from SVG
        function buildEdgeMap() {{
            const edges = document.querySelectorAll('#diagram svg .edge');
            const edgeMap = {{}};  // nodeId -> [edge elements]
            
            edges.forEach(edge => {{
                const title = edge.querySelector('title');
                if (title) {{
                    // Edge title format: "source->target"
                    const parts = title.textContent.split('->');
                    if (parts.length === 2) {{
                        const source = parts[0].trim();
                        const target = parts[1].trim();
                        
                        if (!edgeMap[source]) edgeMap[source] = [];
                        if (!edgeMap[target]) edgeMap[target] = [];
                        
                        edgeMap[source].push({{ edge, type: 'outgoing', target }});
                        edgeMap[target].push({{ edge, type: 'incoming', source }});
                    }}
                }}
            }});
            return edgeMap;
        }}
        
        // Setup hover highlighting on all nodes
        function setupHoverHighlighting() {{
            const diagramEl = document.getElementById('diagram');
            const nodes = document.querySelectorAll('#diagram svg .node');
            const edgeMap = buildEdgeMap();
            let tooltip = null;
            
            nodes.forEach(node => {{
                const titleEl = node.querySelector('title');
                const nodeName = titleEl ? titleEl.textContent : 'Unknown';
                
                node.addEventListener('mouseenter', (e) => {{
                    // Add highlighting class to diagram
                    diagramEl.classList.add('highlighting');
                    
                    // Highlight this node
                    node.classList.add('highlighted');
                    
                    // Highlight connected edges and their nodes
                    const connectedEdges = edgeMap[nodeName] || [];
                    const connectedNodes = new Set();
                    
                    connectedEdges.forEach(({{ edge, type, source, target }}) => {{
                        edge.classList.add('highlighted');
                        if (type === 'outgoing') connectedNodes.add(target);
                        if (type === 'incoming') connectedNodes.add(source);
                    }});
                    
                    // Highlight connected nodes
                    nodes.forEach(n => {{
                        const nTitle = n.querySelector('title');
                        if (nTitle && connectedNodes.has(nTitle.textContent)) {{
                            n.classList.add('highlighted');
                        }}
                    }});
                    
                    // Show tooltip
                    tooltip = document.createElement('div');
                    tooltip.className = 'node-tooltip';
                    tooltip.textContent = nodeName;
                    diagramEl.appendChild(tooltip);
                    
                    const rect = node.getBoundingClientRect();
                    const diagramRect = diagramEl.getBoundingClientRect();
                    tooltip.style.left = (rect.left + rect.width/2 - diagramRect.left + diagramEl.scrollLeft) + 'px';
                    tooltip.style.top = (rect.top - diagramRect.top + diagramEl.scrollTop) + 'px';
                }});
                
                node.addEventListener('mouseleave', () => {{
                    // Remove all highlighting
                    diagramEl.classList.remove('highlighting');
                    document.querySelectorAll('#diagram .highlighted').forEach(el => {{
                        el.classList.remove('highlighted');
                    }});
                    
                    // Remove tooltip
                    if (tooltip) {{
                        tooltip.remove();
                        tooltip = null;
                    }}
                }});
            }});
        }}
        
        function zoomIn() {{ scale *= 1.2; applyZoom(); }}
        function zoomOut() {{ scale *= 0.8; applyZoom(); }}
        function resetZoom() {{ scale = 1; applyZoom(); }}
        
        function applyZoom() {{
            const svg = document.querySelector('#diagram svg');
            if (svg) {{
                svg.style.transform = `scale(${{scale}})`;
                svg.style.transformOrigin = 'top left';
            }}
            // Update zoom level display
            const zoomEl = document.getElementById('zoomLevel');
            if (zoomEl) zoomEl.textContent = Math.round(scale * 100);
        }}
        
        // Drag-to-zoom functionality
        let isSelecting = false;
        let selectionBox = null;
        let startX, startY;
        
        function setupDragToZoom() {{
            const diagramEl = document.getElementById('diagram');
            
            diagramEl.addEventListener('mousedown', (e) => {{
                // Only start selection if clicking on diagram background or SVG
                if (e.target.closest('.node') || e.target.closest('.edge')) return;
                
                isSelecting = true;
                const rect = diagramEl.getBoundingClientRect();
                startX = e.clientX - rect.left + diagramEl.scrollLeft;
                startY = e.clientY - rect.top + diagramEl.scrollTop;
                
                selectionBox = document.createElement('div');
                selectionBox.className = 'selection-box';
                selectionBox.style.left = startX + 'px';
                selectionBox.style.top = startY + 'px';
                diagramEl.appendChild(selectionBox);
            }});
            
            diagramEl.addEventListener('mousemove', (e) => {{
                if (!isSelecting || !selectionBox) return;
                
                const rect = diagramEl.getBoundingClientRect();
                const currentX = e.clientX - rect.left + diagramEl.scrollLeft;
                const currentY = e.clientY - rect.top + diagramEl.scrollTop;
                
                const width = Math.abs(currentX - startX);
                const height = Math.abs(currentY - startY);
                const left = Math.min(startX, currentX);
                const top = Math.min(startY, currentY);
                
                selectionBox.style.left = left + 'px';
                selectionBox.style.top = top + 'px';
                selectionBox.style.width = width + 'px';
                selectionBox.style.height = height + 'px';
            }});
            
            diagramEl.addEventListener('mouseup', (e) => {{
                if (!isSelecting || !selectionBox) return;
                isSelecting = false;
                
                const boxRect = {{
                    width: parseInt(selectionBox.style.width) || 0,
                    height: parseInt(selectionBox.style.height) || 0,
                    left: parseInt(selectionBox.style.left) || 0,
                    top: parseInt(selectionBox.style.top) || 0
                }};
                
                selectionBox.remove();
                selectionBox = null;
                
                // Only zoom if selection is large enough
                if (boxRect.width > 20 && boxRect.height > 20) {{
                    const diagramRect = diagramEl.getBoundingClientRect();
                    const viewWidth = diagramRect.width - 40; // padding
                    const viewHeight = diagramRect.height - 40;
                    
                    // Calculate zoom to fit selection
                    const zoomX = viewWidth / boxRect.width;
                    const zoomY = viewHeight / boxRect.height;
                    const newScale = Math.min(zoomX, zoomY, 5); // cap at 5x
                    
                    scale = newScale;
                    applyZoom();
                    
                    // Scroll to show the selected area
                    diagramEl.scrollLeft = boxRect.left * scale;
                    diagramEl.scrollTop = boxRect.top * scale;
                }}
            }});
            
            diagramEl.addEventListener('mouseleave', () => {{
                if (selectionBox) {{
                    selectionBox.remove();
                    selectionBox = null;
                }}
                isSelecting = false;
            }});
        }}
        
        function downloadSVG() {{
            if (svgContent) {{
                const blob = new Blob([svgContent], {{ type: 'image/svg+xml' }});
                const url = URL.createObjectURL(blob);
                const a = document.createElement('a');
                a.href = url;
                a.download = '{download_name}';
                a.click();
                URL.revokeObjectURL(url);
            }} else {{
                alert('SVG not ready. Please wait for diagram to render.');
            }}
        }}
        
        // Start rendering
        renderDiagram();
    </script>
</body>
</html>'''


def main():
    """Entry point for the CLI."""
    import asyncio
    asyncio.run(run_server())


async def run_server():
    """Run the MCP server on stdio."""
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


if __name__ == "__main__":
    main()
