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

Provides:
- Tool: generate_rtl_image - Generate RTL block diagram image via Nano Banana (Gemini)

Architecture: MCP server sends Verilog source to the Nano Banana intermediate
              server (Vercel or local), which calls Gemini and returns a
              publication-quality image.
"""

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

from rtlviz.telemetry import (
    ping_session_start,
    ping_tool_call,
    ping_diagram_rendered,
)

# Default Nano Banana server URL (set NANOBANANA_URL env var for local dev)
NANOBANANA_URL = os.environ.get("NANOBANANA_URL", "https://rtlviz.vercel.app")

server = Server("rtlviz")

ping_session_start()


@server.list_tools()
async def list_tools() -> list[Tool]:
    """List available tools."""
    return [
        Tool(
            name="generate_rtl_image",
            description="""Generate a professional RTL block diagram IMAGE from Verilog source via Nano Banana (Gemini image generation).

**This tool generates a high-quality raster image (JPEG) of the RTL design,**
suitable for papers, datasheets, and presentations.

**How it works:**
1. Reads all .v files in the source directory
2. Sends source to the Nano Banana server
3. Server calls Gemini image generation (API keys stay server-side)
4. Returns a professional RTL block diagram image

**Features:**
- All multiplexers shown as distinct labeled blocks
- Pipeline stages with clear boundaries
- Signal names with bit widths on all wires
- Color-coded functional grouping
- Publication-quality output""",
            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 image file (will be saved as .jpg)",
                    },
                    "title": {
                        "type": "string",
                        "description": "Title for the diagram",
                        "default": "RTL Block Diagram",
                    },
                    "server_url": {
                        "type": "string",
                        "description": "Nano Banana server URL (default: from NANOBANANA_URL env var)",
                    },
                },
                "required": ["source_dir", "output_path"],
            },
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    """Handle tool calls."""
    ping_tool_call(name)

    if name == "generate_rtl_image":
        return await generate_rtl_image(
            source_dir=arguments["source_dir"],
            output_path=arguments["output_path"],
            title=arguments.get("title", "RTL Block Diagram"),
            server_url=arguments.get("server_url", NANOBANANA_URL),
        )
    return [TextContent(type="text", text=f"Unknown tool: {name}")]


async def generate_rtl_image(
    source_dir: str,
    output_path: str,
    title: str,
    server_url: str = None,
) -> list[TextContent]:
    """Generate an RTL block diagram IMAGE via the Nano Banana intermediate server."""
    if server_url is None:
        server_url = NANOBANANA_URL

    try:
        import urllib.request
        import json

        # Collect Verilog source
        source_path = Path(source_dir)
        v_files = sorted(source_path.rglob("*.v"))
        if not v_files:
            return [TextContent(type="text", text=f"No .v files found in {source_dir}")]

        # Skip testbenches and very large generated files
        skip_patterns = ["_tb.v", "tb_", "testbench"]
        all_source = ""
        included = []
        for vf in v_files:
            name = vf.name.lower()
            if any(p in name for p in skip_patterns):
                continue
            if vf.stat().st_size > 50_000:
                continue
            content = vf.read_text(encoding="utf-8", errors="replace")
            all_source += f"\n// ========== FILE: {vf.name} ==========\n{content}\n"
            included.append(vf.name)

        if not all_source.strip():
            return [TextContent(type="text", text="No valid Verilog source found after filtering.")]

        # Call intermediate server
        payload = json.dumps({
            "verilog_source": all_source,
            "title": title,
        }).encode("utf-8")

        headers = {"Content-Type": "application/json"}
        req = urllib.request.Request(
            f"{server_url}/generate",
            data=payload,
            headers=headers,
            method="POST",
        )

        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                image_bytes = resp.read()
                content_type = resp.headers.get("Content-Type", "image/jpeg")
                time_s = resp.headers.get("X-Time-Seconds", "?")
        except Exception as e:
            error_str = str(e)
            if "Connection refused" in error_str or "URLError" in error_str:
                return [TextContent(
                    type="text",
                    text="Nano Banana server is not running.\n\n"
                         "Start it with:\n"
                         "  python -m gemini_pilot.nanobanana_server\n\n"
                         f"And set NANOBANANA_URL=http://127.0.0.1:8420\n"
                         f"Expected at: {server_url}",
                )]
            return [TextContent(type="text", text=f"Server error: {error_str[:300]}")]

        # Save image
        ext = ".png" if "png" in content_type else ".jpg"
        out_path = Path(output_path).with_suffix(ext)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_bytes(image_bytes)

        ping_diagram_rendered(success=True)

        # Auto-open
        webbrowser.open(f'file:///{out_path.resolve()}')

        size_kb = len(image_bytes) / 1024
        return [TextContent(
            type="text",
            text=f"**Image Generated (Nano Banana)**\n\n"
                 f"- Output: `{out_path}`\n"
                 f"- Size: {size_kb:.0f} KB\n"
                 f"- Time: {time_s}s\n"
                 f"- Files: {len(included)} ({', '.join(included[:5])}{'...' if len(included) > 5 else ''})\n\n"
                 f"Opened in default image viewer.",
        )]
    except Exception as e:
        return [TextContent(type="text", text=f"Error generating RTL image: {e}")]


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()
