#!/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"],
            },
        ),
    ]


@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"),
        )
    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}")]





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, #667eea, #764ba2); 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; }}
        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); }}
        #diagram {{ border: 1px solid #eee; border-radius: 8px; padding: 20px; background: #fafafa; overflow: auto; min-height: 500px; }}
        #diagram svg {{ max-width: 100%; height: auto; }}
        .legend {{ margin-top: 20px; padding: 15px; background: #f8f9fa; border-radius: 8px; display: flex; gap: 25px; flex-wrap: wrap; align-items: center; }}
        .legend-item {{ display: flex; align-items: center; gap: 8px; font-size: 13px; }}
        .legend-line {{ width: 30px; height: 3px; }}
        .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; }}
    </style>
</head>
<body>
    <div class="container">
        <h1>{safe_title}</h1>
        <p class="subtitle">Generated with RTLViz | Diagrams may contain errors - please verify</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>
        </div>
        
        <div id="diagram">
            <div class="loading">Loading Viz.js rendering engine...</div>
        </div>
        
        <div class="legend">
            <strong>Legend:</strong>
            <div class="legend-item">
                <span class="legend-line" style="background: black;"></span>
                <span>Data Path</span>
            </div>
            <div class="legend-item">
                <span class="legend-line" style="background: transparent; border-top: 2px dashed #1565C0;"></span>
                <span>Control Signal</span>
            </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;
            }} 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);
            }}
        }}
        
        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';
            }}
        }}
        
        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()
