#!/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.

**RECOMMENDED WORKFLOW FOR HIGH-QUALITY DIAGRAMS:**
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 }`

**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"
}
```""",
            inputSchema={
                "type": "object",
                "properties": {
                    "dot_content": {
                        "type": "string",
                        "description": "The Graphviz DOT code to render",
                    },
                    "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 BASIC RTL block diagram from Verilog source (pattern-based, minimal detail).

**FOR HIGH-QUALITY DIAGRAMS: Use the Agent-Driven Workflow instead:**
1. Read the `rtlviz://guidelines` resource FIRST
2. Read all .v files in the source directory yourself
3. Parse module structure, registers, clock domains from the code
4. Generate detailed DOT code following the guidelines (white-box detail, all nodes connected)
5. Call `render_diagram` with your DOT code

This tool only provides basic fallback generation.""",
            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["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}")]


async def render_diagram(dot_content: str, output_path: str, title: str) -> list[TextContent]:
    """Generate an interactive HTML viewer from DOT code."""
    try:
        # 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")

        # 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"✅ Successfully generated: {output_path}",
            )
        ]
    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")
        
        # Generate SVG using graphviz Python package
        svg_path = out_path.with_suffix('.svg')
        svg_generated = False
        try:
            import graphviz
            source = graphviz.Source(dot_code)
            svg_content = source.pipe(format='svg').decode('utf-8')
            svg_path.write_text(svg_content, encoding='utf-8')
            svg_generated = True
        except Exception:
            pass  # Graphviz binary not installed
        
        # 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"
        
        svg_info = f"\n- SVG: `{svg_path}`" if svg_generated else "\n- SVG: Not generated (install Graphviz binary)"
        
        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, #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()
