#!/usr/bin/env python3
"""
Enhanced Hanzo MCP Server
========================

MCP server implementation that exposes the 6 universal tools via Model Context Protocol.
Now includes the new unified development tools (edit, fmt, test, build, lint, guard)
plus the existing tool suite.
"""

import json
import asyncio
from typing import Any, Dict, List, Optional, Sequence

from pydantic import AnyUrl
from mcp.types import Tool, Resource, TextContent, ImageContent, LoggingLevel, EmbeddedResource
from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions

from .unified_backend import TargetSpec, ToolResult, backend
from .tools.dev_tools_mcp import dev_tools_server


class HanzoMCPServer:
    """Enhanced MCP server with 6 universal tools"""

    def __init__(self):
        self.server = Server("hanzo-mcp")
        self.dev_tools_server = dev_tools_server  # New orthogonal dev tools
        self.setup_handlers()

    def setup_handlers(self):
        """Set up MCP server handlers"""

        @self.server.list_tools()
        async def handle_list_tools() -> List[Tool]:
            """List available tools"""
            return [
                Tool(
                    name="edit",
                    description="Semantic refactors via LSP across languages",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "target": {
                                "type": "string",
                                "description": "Target: file:<path>, dir:<path>, pkg:<spec>, ws, or changed",
                            },
                            "op": {
                                "type": "string",
                                "enum": ["rename", "code_action", "organize_imports", "apply_workspace_edit"],
                                "description": "Operation to perform",
                            },
                            "language": {
                                "type": "string",
                                "default": "auto",
                                "description": "Language override (auto, go, ts, py, rs, cc, sol)",
                            },
                            "backend": {"type": "string", "default": "auto", "description": "Backend override"},
                            "root": {"type": "string", "description": "Workspace root override"},
                            "env": {
                                "type": "object",
                                "additionalProperties": {"type": "string"},
                                "description": "Environment variables",
                            },
                            "file": {"type": "string", "description": "File path for rename/code_action operations"},
                            "pos": {
                                "type": "object",
                                "properties": {"line": {"type": "integer"}, "character": {"type": "integer"}},
                                "description": "Position for rename/code_action",
                            },
                            "range": {
                                "type": "object",
                                "properties": {
                                    "start": {
                                        "type": "object",
                                        "properties": {"line": {"type": "integer"}, "character": {"type": "integer"}},
                                    },
                                    "end": {
                                        "type": "object",
                                        "properties": {"line": {"type": "integer"}, "character": {"type": "integer"}},
                                    },
                                },
                                "description": "Range for code actions",
                            },
                            "new_name": {"type": "string", "description": "New name for rename operation"},
                            "only": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "Code action kinds filter",
                            },
                            "apply": {
                                "type": "boolean",
                                "default": True,
                                "description": "Apply changes (default true)",
                            },
                            "workspace_edit": {
                                "type": "object",
                                "description": "WorkspaceEdit payload for apply_workspace_edit",
                            },
                            "dry_run": {
                                "type": "boolean",
                                "default": False,
                                "description": "Preview mode - no file writes",
                            },
                        },
                        "required": ["target", "op"],
                        "additionalProperties": False,
                    },
                ),
                Tool(
                    name="fmt",
                    description="Formatting + import normalization",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "target": {
                                "type": "string",
                                "description": "Target: file:<path>, dir:<path>, pkg:<spec>, ws, or changed",
                            },
                            "language": {"type": "string", "default": "auto", "description": "Language override"},
                            "backend": {"type": "string", "default": "auto", "description": "Backend override"},
                            "root": {"type": "string", "description": "Workspace root override"},
                            "env": {
                                "type": "object",
                                "additionalProperties": {"type": "string"},
                                "description": "Environment variables",
                            },
                            "opts": {
                                "type": "object",
                                "properties": {
                                    "local_prefix": {
                                        "type": "string",
                                        "description": "Go import grouping prefix (e.g. github.com/luxfi)",
                                    }
                                },
                                "description": "Tool-specific options",
                            },
                            "dry_run": {"type": "boolean", "default": False, "description": "Preview mode"},
                        },
                        "required": ["target"],
                        "additionalProperties": False,
                    },
                ),
                Tool(
                    name="test",
                    description="Run tests narrowly by default",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "target": {
                                "type": "string",
                                "description": "Target: file:<path>, dir:<path>, pkg:<spec>, ws, or changed",
                            },
                            "language": {"type": "string", "default": "auto", "description": "Language override"},
                            "backend": {"type": "string", "default": "auto", "description": "Backend override"},
                            "root": {"type": "string", "description": "Workspace root override"},
                            "env": {
                                "type": "object",
                                "additionalProperties": {"type": "string"},
                                "description": "Environment variables",
                            },
                            "opts": {
                                "type": "object",
                                "properties": {
                                    "run_filter": {"type": "string", "description": "Test filter pattern"},
                                    "count": {"type": "integer", "description": "Number of times to run each test"},
                                    "race": {"type": "boolean", "description": "Enable race detection (Go)"},
                                    "watch": {"type": "boolean", "default": False, "description": "Watch mode"},
                                },
                                "description": "Test-specific options",
                            },
                            "dry_run": {"type": "boolean", "default": False, "description": "Preview mode"},
                        },
                        "required": ["target"],
                        "additionalProperties": False,
                    },
                ),
                Tool(
                    name="build",
                    description="Compile/build artifacts narrowly by default",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "target": {
                                "type": "string",
                                "description": "Target: file:<path>, dir:<path>, pkg:<spec>, ws, or changed",
                            },
                            "language": {"type": "string", "default": "auto", "description": "Language override"},
                            "backend": {"type": "string", "default": "auto", "description": "Backend override"},
                            "root": {"type": "string", "description": "Workspace root override"},
                            "env": {
                                "type": "object",
                                "additionalProperties": {"type": "string"},
                                "description": "Environment variables",
                            },
                            "opts": {
                                "type": "object",
                                "properties": {
                                    "release": {"type": "boolean", "default": False, "description": "Release build"},
                                    "features": {
                                        "type": "array",
                                        "items": {"type": "string"},
                                        "description": "Rust features to enable",
                                    },
                                },
                                "description": "Build-specific options",
                            },
                            "dry_run": {"type": "boolean", "default": False, "description": "Preview mode"},
                        },
                        "required": ["target"],
                        "additionalProperties": False,
                    },
                ),
                Tool(
                    name="lint",
                    description="Lint/typecheck in one place",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "target": {
                                "type": "string",
                                "description": "Target: file:<path>, dir:<path>, pkg:<spec>, ws, or changed",
                            },
                            "language": {"type": "string", "default": "auto", "description": "Language override"},
                            "backend": {"type": "string", "default": "auto", "description": "Backend override"},
                            "root": {"type": "string", "description": "Workspace root override"},
                            "env": {
                                "type": "object",
                                "additionalProperties": {"type": "string"},
                                "description": "Environment variables",
                            },
                            "opts": {
                                "type": "object",
                                "properties": {
                                    "fix": {
                                        "type": "boolean",
                                        "default": False,
                                        "description": "Auto-fix issues where possible",
                                    }
                                },
                                "description": "Lint-specific options",
                            },
                            "dry_run": {"type": "boolean", "default": False, "description": "Preview mode"},
                        },
                        "required": ["target"],
                        "additionalProperties": False,
                    },
                ),
                Tool(
                    name="guard",
                    description="Repo invariants (boundaries, forbidden imports/strings)",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "target": {
                                "type": "string",
                                "description": "Target: file:<path>, dir:<path>, pkg:<spec>, ws, or changed",
                            },
                            "language": {"type": "string", "default": "auto", "description": "Language override"},
                            "backend": {"type": "string", "default": "auto", "description": "Backend override"},
                            "root": {"type": "string", "description": "Workspace root override"},
                            "env": {
                                "type": "object",
                                "additionalProperties": {"type": "string"},
                                "description": "Environment variables",
                            },
                            "rules": {
                                "type": "array",
                                "items": {
                                    "type": "object",
                                    "properties": {
                                        "id": {"type": "string"},
                                        "type": {"type": "string", "enum": ["regex", "import", "generated"]},
                                        "glob": {"type": "string"},
                                        "pattern": {"type": "string"},
                                        "forbid_import_prefix": {"type": "string"},
                                        "forbid_writes": {"type": "boolean"},
                                    },
                                    "required": ["id", "type", "glob"],
                                },
                                "description": "Guard rules to check",
                            },
                            "dry_run": {"type": "boolean", "default": False, "description": "Preview mode"},
                        },
                        "required": ["target", "rules"],
                        "additionalProperties": False,
                    },
                ),
                Tool(
                    name="search_codebase",
                    description="Search for symbols, files, or content across codebase",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Search query (symbol name, file pattern, or content)",
                            },
                            "type": {
                                "type": "string",
                                "enum": ["symbols", "files", "content"],
                                "default": "symbols",
                                "description": "Search type",
                            },
                            "language": {"type": "string", "description": "Filter by language"},
                            "limit": {"type": "integer", "default": 50, "description": "Maximum results"},
                        },
                        "required": ["query"],
                    },
                ),
                Tool(
                    name="get_session_history",
                    description="Get recent MCP tool usage history",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "limit": {"type": "integer", "default": 10, "description": "Number of recent sessions"}
                        },
                    },
                ),
            ]

        @self.server.call_tool()
        async def handle_call_tool(name: str, arguments: dict) -> List[TextContent]:
            """Handle tool calls"""
            try:
                if name == "edit":
                    target_spec = TargetSpec(
                        target=arguments["target"],
                        language=arguments.get("language", "auto"),
                        backend=arguments.get("backend", "auto"),
                        root=arguments.get("root"),
                        env=arguments.get("env", {}),
                        dry_run=arguments.get("dry_run", False),
                    )

                    result = await backend.edit(
                        target_spec,
                        op=arguments["op"],
                        file=arguments.get("file"),
                        pos=arguments.get("pos"),
                        range=arguments.get("range"),
                        new_name=arguments.get("new_name"),
                        only=arguments.get("only", []),
                    )

                    backend.session_manager.log_tool_execution("edit", arguments, result)

                    return [TextContent(type="text", text=self._format_tool_result(result))]

                elif name == "fmt":
                    target_spec = TargetSpec(
                        target=arguments["target"],
                        language=arguments.get("language", "auto"),
                        backend=arguments.get("backend", "auto"),
                        root=arguments.get("root"),
                        env=arguments.get("env", {}),
                        dry_run=arguments.get("dry_run", False),
                    )

                    result = await backend.fmt(target_spec, opts=arguments.get("opts", {}))

                    backend.session_manager.log_tool_execution("fmt", arguments, result)

                    return [TextContent(type="text", text=self._format_tool_result(result))]

                elif name == "test":
                    target_spec = TargetSpec(
                        target=arguments["target"],
                        language=arguments.get("language", "auto"),
                        backend=arguments.get("backend", "auto"),
                        root=arguments.get("root"),
                        env=arguments.get("env", {}),
                        dry_run=arguments.get("dry_run", False),
                    )

                    result = await backend.test(target_spec, opts=arguments.get("opts", {}))

                    backend.session_manager.log_tool_execution("test", arguments, result)

                    return [TextContent(type="text", text=self._format_tool_result(result))]

                elif name == "build":
                    target_spec = TargetSpec(
                        target=arguments["target"],
                        language=arguments.get("language", "auto"),
                        backend=arguments.get("backend", "auto"),
                        root=arguments.get("root"),
                        env=arguments.get("env", {}),
                        dry_run=arguments.get("dry_run", False),
                    )

                    result = await backend.build(target_spec, opts=arguments.get("opts", {}))

                    backend.session_manager.log_tool_execution("build", arguments, result)

                    return [TextContent(type="text", text=self._format_tool_result(result))]

                elif name == "lint":
                    target_spec = TargetSpec(
                        target=arguments["target"],
                        language=arguments.get("language", "auto"),
                        backend=arguments.get("backend", "auto"),
                        root=arguments.get("root"),
                        env=arguments.get("env", {}),
                        dry_run=arguments.get("dry_run", False),
                    )

                    result = await backend.lint(target_spec, opts=arguments.get("opts", {}))

                    backend.session_manager.log_tool_execution("lint", arguments, result)

                    return [TextContent(type="text", text=self._format_tool_result(result))]

                elif name == "guard":
                    target_spec = TargetSpec(
                        target=arguments["target"],
                        language=arguments.get("language", "auto"),
                        backend=arguments.get("backend", "auto"),
                        root=arguments.get("root"),
                        env=arguments.get("env", {}),
                        dry_run=arguments.get("dry_run", False),
                    )

                    result = await backend.guard(target_spec, rules=arguments["rules"])

                    backend.session_manager.log_tool_execution("guard", arguments, result)

                    return [TextContent(type="text", text=self._format_tool_result(result))]

                elif name == "search_codebase":
                    query = arguments["query"]
                    search_type = arguments.get("type", "symbols")
                    language = arguments.get("language")
                    limit = arguments.get("limit", 50)

                    if search_type == "symbols":
                        results = backend.indexer.search_symbols(query, language)
                        results_text = "\n".join(
                            [f"{r['path']}:{r['line']} - {r['kind']} {r['name']}" for r in results[:limit]]
                        )
                    else:
                        results_text = f"Search type '{search_type}' not implemented yet"

                    return [TextContent(type="text", text=f"Search results for '{query}':\n{results_text}")]

                elif name == "get_session_history":
                    limit = arguments.get("limit", 10)
                    sessions = backend.session_manager.get_recent_sessions(limit)

                    history_text = "Recent MCP Sessions:\n"
                    for session in sessions:
                        history_text += f"\nSession {session['session_id'][:8]}... ({session['start_time']})\n"
                        history_text += f"  Tools: {', '.join(session['tools_used'])} ({session['tool_count']} calls)\n"

                    return [TextContent(type="text", text=history_text)]

                else:
                    return [TextContent(type="text", text=f"Unknown tool: {name}")]

            except Exception as e:
                return [TextContent(type="text", text=f"Error executing {name}: {str(e)}")]

    def _format_tool_result(self, result: ToolResult) -> str:
        """Format tool result for display"""
        status = "✅ SUCCESS" if result.ok else "❌ FAILED"

        output = f"{status} ({result.execution_time:.2f}s)\n"
        output += f"Root: {result.root}\n"
        output += f"Language: {result.language_used}\n"
        output += f"Backend: {result.backend_used}\n"

        if result.scope_resolved:
            output += f"Scope: {len(result.scope_resolved)} files\n"

        if result.touched_files:
            output += f"Modified: {len(result.touched_files)} files\n"
            for f in result.touched_files[:5]:  # Show first 5
                output += f"  - {f}\n"
            if len(result.touched_files) > 5:
                output += f"  ... and {len(result.touched_files) - 5} more\n"

        if result.stdout:
            output += f"\nOutput:\n{result.stdout}\n"

        if result.stderr:
            output += f"\nErrors:\n{result.stderr}\n"

        if result.errors:
            output += f"\nIssues:\n"
            for error in result.errors:
                output += f"  - {error}\n"

        return output

    async def run(self, transport_type: str = "stdio"):
        """Run the MCP server"""
        if transport_type == "stdio":
            from mcp.server.stdio import stdio_server

            async with stdio_server() as (read_stream, write_stream):
                await self.server.run(
                    read_stream,
                    write_stream,
                    InitializationOptions(
                        server_name="hanzo-mcp",
                        server_version="1.0.0",
                        capabilities=self.server.get_capabilities(
                            notification_options=NotificationOptions(), experimental_capabilities={}
                        ),
                    ),
                )


if __name__ == "__main__":
    server = HanzoMCPServer()
    asyncio.run(server.run())
