"""NAVIG MCP Server

Model Context Protocol server that exposes NAVIG's capabilities to AI assistants.
This allows Copilot, Claude, and other MCP-compatible assistants to:
- Query hosts, apps, and database configurations
- Search the wiki knowledge base
- Execute commands on remote servers
- Access project context

Usage:
    navig mcp serve                          # Start MCP server (stdio mode)
    navig mcp serve --transport websocket    # Start WebSocket server on port 3001
    navig mcp serve --port 3001             # Same as above (port implies websocket)

VS Code Integration:
    Add to .vscode/mcp.json or VS Code settings
"""

import asyncio
import json
import secrets
import sys
from collections.abc import Callable
from datetime import datetime
from pathlib import Path
from typing import Any

from navig import console_helper as ch
from navig.config import ConfigManager


class MCPProtocolHandler:
    """Handles MCP JSON-RPC protocol over stdio."""

    def __init__(self):
        self.tools: dict[str, dict[str, Any]] = {}
        self.resources: dict[str, dict[str, Any]] = {}
        self.prompts: dict[str, dict[str, Any]] = {}
        self._handlers: dict[str, Callable] = {}
        self._running = False
        self._config = ConfigManager()

        # Register built-in methods
        self._setup_protocol_handlers()
        self._setup_navig_tools()
        self._setup_navig_resources()

    def _setup_protocol_handlers(self):
        """Setup MCP protocol method handlers."""
        self._handlers = {}
        self._handlers["initialize"] = self._handle_initialize
        self._handlers["initialized"] = self._handle_initialized
        self._handlers["tools/list"] = self._handle_tools_list
        self._handlers["tools/call"] = self._handle_tools_call
        self._handlers["resources/list"] = self._handle_resources_list
        self._handlers["resources/read"] = self._handle_resources_read
        self._handlers["prompts/list"] = self._handle_prompts_list
        self._handlers["prompts/get"] = self._handle_prompts_get
        self._handlers["ping"] = self._handle_ping

    def _setup_navig_tools(self):
        """Register NAVIG tools for MCP."""
        self.tools = {}
        self._tool_handlers = {}
        from navig.mcp.tools import register_all_tools

        register_all_tools(self)

    def _setup_navig_resources(self):
        """Register NAVIG resources for MCP."""
        self.resources = {
            "navig://config/hosts": {
                "uri": "navig://config/hosts",
                "name": "NAVIG Hosts Configuration",
                "description": "All configured SSH hosts",
                "mimeType": "application/json",
            },
            "navig://config/apps": {
                "uri": "navig://config/apps",
                "name": "NAVIG Apps Configuration",
                "description": "All configured applications",
                "mimeType": "application/json",
            },
            "navig://wiki": {
                "uri": "navig://wiki",
                "name": "NAVIG Wiki",
                "description": "Project knowledge base",
                "mimeType": "text/markdown",
            },
            "navig://context": {
                "uri": "navig://context",
                "name": "NAVIG Context",
                "description": "Current system state and recent errors",
                "mimeType": "application/json",
            },
            "navig://agent/status": {
                "uri": "navig://agent/status",
                "name": "NAVIG Agent Status",
                "description": "Agent runtime status and configuration",
                "mimeType": "application/json",
            },
            "navig://agent/goals": {
                "uri": "navig://agent/goals",
                "name": "NAVIG Agent Goals",
                "description": "Autonomous goal list and progress",
                "mimeType": "application/json",
            },
            "navig://agent/remediation": {
                "uri": "navig://agent/remediation",
                "name": "NAVIG Agent Remediation",
                "description": "Remediation actions and recent remediation logs",
                "mimeType": "application/json",
            },
            "navig://agent/learning": {
                "uri": "navig://agent/learning",
                "name": "NAVIG Agent Learning Report",
                "description": "Latest error pattern analysis report",
                "mimeType": "application/json",
            },
            "navig://agent/service": {
                "uri": "navig://agent/service",
                "name": "NAVIG Agent Service Status",
                "description": "Service installer/status integration",
                "mimeType": "application/json",
            },
            "agent://status": {
                "uri": "agent://status",
                "name": "Agent Status",
                "description": "Alias for NAVIG agent runtime status",
                "mimeType": "application/json",
            },
            "agent://goals": {
                "uri": "agent://goals",
                "name": "Agent Goals",
                "description": "Alias for agent goals",
                "mimeType": "application/json",
            },
            "agent://remediation": {
                "uri": "agent://remediation",
                "name": "Agent Remediation",
                "description": "Alias for remediation actions",
                "mimeType": "application/json",
            },
            "agent://learning/patterns": {
                "uri": "agent://learning/patterns",
                "name": "Agent Learning Patterns",
                "description": "Alias for learning report",
                "mimeType": "application/json",
            },
            "agent://service": {
                "uri": "agent://service",
                "name": "Agent Service",
                "description": "Alias for service status",
                "mimeType": "application/json",
            },
            # ── Runtime Contracts ──────────────────────────────────────
            "navig://runtime/nodes": {
                "uri": "navig://runtime/nodes",
                "name": "NAVIG Runtime Nodes",
                "description": "All registered Node identities",
                "mimeType": "application/json",
            },
            "navig://runtime/missions": {
                "uri": "navig://runtime/missions",
                "name": "NAVIG Runtime Missions",
                "description": "Recent Missions with lifecycle state",
                "mimeType": "application/json",
            },
            "navig://runtime/receipts": {
                "uri": "navig://runtime/receipts",
                "name": "NAVIG Execution Receipts",
                "description": "Audit trail of completed Mission executions",
                "mimeType": "application/json",
            },
        }

    # =========================================================================
    # MCP Protocol Handlers
    # =========================================================================

    def _handle_initialize(self, params: dict[str, Any]) -> dict[str, Any]:
        """Handle MCP initialize request."""
        return {
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {}, "resources": {}, "prompts": {}},
            "serverInfo": {"name": "navig-mcp-server", "version": "1.0.0"},
        }

    def _handle_initialized(self, params: dict[str, Any]) -> None:
        """Handle MCP initialized notification."""
        return None

    def _handle_ping(self, params: dict[str, Any]) -> dict[str, Any]:
        """Handle ping request."""
        return {}

    def _handle_tools_list(self, params: dict[str, Any]) -> dict[str, Any]:
        """Return list of available tools."""
        return {"tools": list(self.tools.values())}

    def _handle_tools_call(self, params: dict[str, Any]) -> dict[str, Any]:
        """Execute a tool call."""
        tool_name = params.get("name")
        arguments = params.get("arguments", {})

        if tool_name not in self.tools:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Unknown tool: {tool_name}"}],
            }

        try:
            result = self._execute_tool(tool_name, arguments)
            return {
                "content": [{"type": "text", "text": json.dumps(result, indent=2, default=str)}]
            }
        except Exception as e:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Tool error: {str(e)}"}],
            }

    def _handle_resources_list(self, params: dict[str, Any]) -> dict[str, Any]:
        """Return list of available resources."""
        return {"resources": list(self.resources.values())}

    def _handle_resources_read(self, params: dict[str, Any]) -> dict[str, Any]:
        """Read a resource."""
        uri = params.get("uri")

        if uri not in self.resources:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Unknown resource: {uri}"}],
            }

        try:
            content = self._read_resource(uri)
            return {
                "contents": [
                    {
                        "uri": uri,
                        "mimeType": self.resources[uri].get("mimeType", "text/plain"),
                        "text": content,
                    }
                ]
            }
        except Exception as e:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Resource error: {str(e)}"}],
            }

    def _handle_prompts_list(self, params: dict[str, Any]) -> dict[str, Any]:
        """Return list of available prompts."""
        return {"prompts": []}

    def _handle_prompts_get(self, params: dict[str, Any]) -> dict[str, Any]:
        """Get a specific prompt."""
        return {"messages": []}

    # =========================================================================
    # Tool Implementations
    # =========================================================================

    def _execute_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any:
        """Execute a tool and return result."""
        handler = getattr(self, "_tool_handlers", {}).get(tool_name)
        if handler:
            return handler(self, arguments)
        raise ValueError(f"Unknown tool: {tool_name}")

    def _read_resource(self, uri: str) -> str:
        """Read resource content."""
        if uri == "navig://config/hosts":
            return json.dumps(self._execute_tool("navig_list_hosts", {}), indent=2)
        elif uri == "navig://config/apps":
            return json.dumps(self._execute_tool("navig_list_apps", {}), indent=2)
        elif uri == "navig://wiki":
            pages = self._execute_tool("navig_wiki_list", {})
            if isinstance(pages, dict) and "error" in pages:
                return f"# Wiki\n\n{pages['error']}"
            return "# Wiki Index\n\n" + "\n".join(f"- [{p['title']}]({p['path']})" for p in pages)
        elif uri == "navig://context":
            return json.dumps(self._execute_tool("navig_get_context", {}), indent=2)
        elif uri in ("navig://agent/status", "agent://status"):
            return json.dumps(self._execute_tool("navig_agent_status_get", {}), indent=2)
        elif uri in ("navig://agent/goals", "agent://goals"):
            return json.dumps(self._execute_tool("navig_agent_goal_list", {"limit": 100}), indent=2)
        elif uri in ("navig://agent/remediation", "agent://remediation"):
            return json.dumps(
                self._execute_tool("navig_agent_remediation_list", {"limit": 100}),
                indent=2,
            )
        elif uri in ("navig://agent/learning", "agent://learning/patterns"):
            report_path = Path.home() / ".navig" / "workspace" / "error-patterns.json"
            if report_path.exists():
                return report_path.read_text(encoding="utf-8", errors="replace")
            return json.dumps(
                self._execute_tool("navig_agent_learning_run", {"days": 7, "export": False}),
                indent=2,
            )
        elif uri in ("navig://agent/service", "agent://service"):
            return json.dumps(self._execute_tool("navig_agent_service_status", {}), indent=2)
        elif uri == "navig://runtime/nodes":
            from navig.contracts.store import get_runtime_store

            store = get_runtime_store()
            return json.dumps([n.to_dict() for n in store.list_nodes()], indent=2)
        elif uri == "navig://runtime/missions":
            from navig.contracts.store import get_runtime_store

            store = get_runtime_store()
            return json.dumps([m.to_dict() for m in store.list_missions(limit=50)], indent=2)
        elif uri == "navig://runtime/receipts":
            from navig.contracts.store import get_runtime_store

            store = get_runtime_store()
            return json.dumps([r.to_dict() for r in store.list_receipts(limit=50)], indent=2)
        else:
            raise ValueError(f"Unknown resource: {uri}")

    # =========================================================================
    # Server Loop
    # =========================================================================

    def handle_message(self, message: dict[str, Any]) -> dict[str, Any] | None:
        """Handle a single JSON-RPC message."""
        method = message.get("method")
        params = message.get("params", {})
        msg_id = message.get("id")

        if method not in self._handlers:
            if msg_id is not None:
                return {
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "error": {"code": -32601, "message": f"Method not found: {method}"},
                }
            return None

        try:
            result = self._handlers[method](params)
            if msg_id is not None:
                return {"jsonrpc": "2.0", "id": msg_id, "result": result}
            return None
        except Exception as e:
            if msg_id is not None:
                return {
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "error": {"code": -32603, "message": str(e)},
                }
            return None

    def run_stdio(self):
        """Run MCP server in stdio mode."""
        self._running = True

        while self._running:
            try:
                line = sys.stdin.readline()
                if not line:
                    break

                message = json.loads(line)
                response = self.handle_message(message)

                if response:
                    sys.stdout.write(json.dumps(response) + "\n")
                    sys.stdout.flush()

            except json.JSONDecodeError as e:
                sys.stderr.write(f"JSON decode error: {e}\n")
            except Exception as e:
                sys.stderr.write(f"Error: {e}\n")

    def stop(self):
        """Stop the server."""
        self._running = False


def start_mcp_server(mode: str = "stdio", port: int = 3001, token: str | None = None):
    """Start the NAVIG MCP server.

    Args:
        mode: Server mode - 'stdio' for stdin/stdout, 'websocket' for WS
        port: Port for WebSocket mode (default 3001)
        token: Optional auth token; auto-generated if None in websocket mode
    """
    handler = MCPProtocolHandler()

    if mode == "stdio":
        handler.run_stdio()
    elif mode == "websocket":
        _run_websocket_server(handler, port, token)
    else:
        raise ValueError(f"Unknown mode: {mode}. Use 'stdio' or 'websocket'.")


def _run_websocket_server(handler: MCPProtocolHandler, port: int, token: str | None = None):
    """Start a WebSocket MCP server with optional token auth.

    The server speaks JSON-RPC 2.0 over WebSocket frames.
    Each frame is one JSON-RPC request/response.
    Notifications (no 'id') are fire-and-forget.

    Auth: if token is set, clients must send it as the first message
    or via the ``Authorization`` header (``Bearer <token>``).
    """
    try:
        import websockets
        import websockets.asyncio.server
    except ImportError as _exc:
        ch.error("WebSocket mode requires the 'websockets' package.")
        ch.info("Install with:  pip install websockets")
        raise SystemExit(1) from _exc

    # Generate a session token if not provided
    session_token = token or secrets.token_urlsafe(32)
    authenticated_clients: set = set()
    authenticated_websockets: set = set()

    # --- Event Bridge integration (push pipeline) ---
    from navig.event_bridge import EventBridge, SubscriptionFilter

    event_bridge = EventBridge(debounce_seconds=1.0)

    notification_sources = {
        "agent.status.changed": "navig://agent/status",
        "agent.goal.changed": "navig://agent/goals",
        "agent.remediation.changed": "navig://agent/remediation",
    }
    debounce_seconds = 1.5
    poll_interval_seconds = 1.0
    max_notification_bytes = 131072
    notification_state: dict[str, dict[str, Any]] = {
        topic: {
            "last_seen_digest": None,
            "last_seen_at": None,
            "last_payload": None,
            "last_emitted_digest": None,
        }
        for topic in notification_sources
    }

    async def _broadcast_notification(payload: str) -> None:
        """Broadcast notification to authenticated clients with drop-safe sends."""
        if not authenticated_websockets:
            return

        async def _send_safe(websocket_obj) -> bool:
            try:
                await asyncio.wait_for(websocket_obj.send(payload), timeout=0.25)
                return True
            except Exception:
                return False

        sockets = list(authenticated_websockets)
        results = await asyncio.gather(
            *[_send_safe(ws) for ws in sockets],
            return_exceptions=True,
        )
        for ws, result in zip(sockets, results):
            if result is not True:
                authenticated_websockets.discard(ws)

    async def _notification_loop() -> None:
        """Poll runtime resources and emit debounced MCP notifications on change."""
        import hashlib

        while True:
            await asyncio.sleep(poll_interval_seconds)
            if not authenticated_websockets:
                continue

            now = datetime.now()
            for topic, uri in notification_sources.items():
                try:
                    resource_text = handler._read_resource(uri)
                except Exception:
                    continue

                digest = hashlib.sha256(resource_text.encode("utf-8", errors="replace")).hexdigest()

                state = notification_state[topic]
                if digest != state["last_seen_digest"]:
                    state["last_seen_digest"] = digest
                    state["last_seen_at"] = now

                    params: dict[str, Any] = {
                        "uri": uri,
                        "changed_at": now.isoformat(),
                    }
                    try:
                        params["data"] = json.loads(resource_text)
                    except json.JSONDecodeError:
                        params["text"] = resource_text[:4000]

                    payload_obj = {
                        "jsonrpc": "2.0",
                        "method": topic,
                        "params": params,
                    }
                    payload = json.dumps(payload_obj, default=str)
                    if len(payload.encode("utf-8", errors="replace")) > max_notification_bytes:
                        payload = json.dumps(
                            {
                                "jsonrpc": "2.0",
                                "method": topic,
                                "params": {
                                    "uri": uri,
                                    "changed_at": now.isoformat(),
                                    "truncated": True,
                                },
                            },
                            default=str,
                        )
                    state["last_payload"] = payload

                # Debounce: only emit after value has stayed unchanged for a short window.
                if (
                    state["last_seen_digest"] is not None
                    and state["last_seen_digest"] != state["last_emitted_digest"]
                    and state["last_seen_at"] is not None
                    and (now - state["last_seen_at"]).total_seconds() >= debounce_seconds
                    and state["last_payload"] is not None
                ):
                    await _broadcast_notification(state["last_payload"])
                    # Also push through EventBridge for filtered delivery
                    from navig.event_bridge import Severity

                    await event_bridge.push_direct(
                        topic=topic,
                        source="mcp.resource_poll",
                        data={"uri": uri, "changed_at": now.isoformat()},
                        severity=Severity.INFO,
                    )
                    state["last_emitted_digest"] = state["last_seen_digest"]

    async def _handle_client(websocket):
        client_id = id(websocket)

        # Check auth header first
        auth_header = websocket.request.headers.get("Authorization", "")
        if auth_header == f"Bearer {session_token}":
            authenticated_clients.add(client_id)
            authenticated_websockets.add(websocket)
            event_bridge.register_client(websocket)

        try:
            async for raw in websocket:
                # If not yet authenticated, expect token as first message
                if session_token and client_id not in authenticated_clients:
                    if raw.strip() == session_token:
                        authenticated_clients.add(client_id)
                        authenticated_websockets.add(websocket)
                        event_bridge.register_client(websocket)
                        ack = json.dumps(
                            {
                                "jsonrpc": "2.0",
                                "result": {"authenticated": True},
                                "id": 0,
                            }
                        )
                        await websocket.send(ack)
                        continue
                    else:
                        err = json.dumps(
                            {
                                "jsonrpc": "2.0",
                                "error": {
                                    "code": -32000,
                                    "message": "Authentication required",
                                },
                                "id": None,
                            }
                        )
                        await websocket.send(err)
                        await websocket.close(4001, "Authentication failed")
                        return

                # Normal JSON-RPC handling
                try:
                    message = json.loads(raw)
                except json.JSONDecodeError:
                    err = json.dumps(
                        {
                            "jsonrpc": "2.0",
                            "error": {"code": -32700, "message": "Parse error"},
                            "id": None,
                        }
                    )
                    await websocket.send(err)
                    continue

                response = handler.handle_message(message)

                # --- EventBridge: handle subscription requests ---
                method = message.get("method", "")
                msg_id = message.get("id")

                if method == "navig.subscribe" and msg_id is not None:
                    params = message.get("params", {})
                    filt = SubscriptionFilter.from_dict(params)
                    event_bridge.update_client_filter(websocket, filt)
                    sub_resp = {
                        "jsonrpc": "2.0",
                        "result": {"subscribed": True, "filter": filt.to_dict()},
                        "id": msg_id,
                    }
                    await websocket.send(json.dumps(sub_resp, default=str))
                    continue

                if method == "navig.bridge.stats" and msg_id is not None:
                    stats_resp = {
                        "jsonrpc": "2.0",
                        "result": event_bridge.get_stats(),
                        "id": msg_id,
                    }
                    await websocket.send(json.dumps(stats_resp, default=str))
                    continue

                if response:
                    await websocket.send(json.dumps(response, default=str))
        except websockets.exceptions.ConnectionClosed:
            pass  # connection closed normally
        finally:
            authenticated_clients.discard(client_id)
            authenticated_websockets.discard(websocket)
            event_bridge.unregister_client(websocket)

    async def _serve():
        import contextlib

        ch.success(f"NAVIG MCP WebSocket server listening on ws://localhost:{port}")
        ch.info(f"Session token: {session_token}")
        ch.dim("Clients must authenticate with the token before sending requests.")
        ch.dim("Press Ctrl+C to stop.")

        notifier_task = asyncio.create_task(_notification_loop())
        async with websockets.asyncio.server.serve(
            _handle_client,
            "0.0.0.0",
            port,
            ping_interval=30,
            ping_timeout=10,
        ):
            try:
                await asyncio.Future()  # run forever
            finally:
                notifier_task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await notifier_task

    try:
        asyncio.run(_serve())
    except KeyboardInterrupt:
        ch.info("MCP WebSocket server stopped.")


def generate_vscode_mcp_config() -> dict[str, Any]:
    """Generate VS Code MCP configuration for Copilot integration."""
    import sys

    python_path = sys.executable

    return {
        "mcpServers": {
            "navig": {
                "command": python_path,
                "args": ["-m", "navig.mcp_server"],
                "env": {},
            }
        }
    }


def generate_claude_mcp_config() -> dict[str, Any]:
    """Generate Claude Desktop MCP configuration."""
    import sys

    python_path = sys.executable

    return {
        "mcpServers": {
            "navig": {
                "command": python_path,
                "args": ["-m", "navig.mcp_server"],
                "env": {},
            }
        }
    }


# =============================================================================
# Memory MCP Tool Handlers (module-level, importable for testing)
# =============================================================================

from navig.memory.paths import KEY_FACTS_DB_PATH as _KEY_FACTS_DB_PATH


def _memory_store():
    """Return a fresh KeyFactStore backed by the canonical DB path."""
    from navig.memory.key_facts import KeyFactStore

    return KeyFactStore(db_path=_KEY_FACTS_DB_PATH)


async def memory_retrieve(query: str, limit: int = 10, token_budget: int = 2000) -> dict:
    """Retrieve ranked key facts matching query within token budget."""
    from navig.memory.fact_retriever import FactRetriever

    retriever = FactRetriever(_memory_store())
    facts = retriever.retrieve(query=query, limit=limit, token_budget=token_budget)
    return {"facts": [f.model_dump() if hasattr(f, "model_dump") else vars(f) for f in facts]}


async def memory_remember(text: str, source: str = "mcp") -> dict:
    """Extract and persist key facts from text."""
    from navig.memory.fact_extractor import FactExtractor

    store = _memory_store()
    extractor = FactExtractor(store)
    added = await extractor.extract_and_store(
        user_text=text, assistant_text="", source_platform=source
    )
    return {"added": added}


async def memory_forget(fact_id: str) -> dict:
    """Soft-delete a key fact by ID."""
    store = _memory_store()
    ok = store.soft_delete(fact_id)
    return {"deleted": ok, "id": fact_id}


async def memory_stats() -> dict:
    """Return key fact store statistics."""
    store = _memory_store()
    return store.stats()


# Allow running as module: python -m navig.mcp_server
# Usage:
#   python -m navig.mcp_server                  # stdio mode
#   python -m navig.mcp_server --websocket      # WebSocket on port 3001
#   python -m navig.mcp_server --port 3001      # Same as above
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="NAVIG MCP Server")
    parser.add_argument("--websocket", action="store_true", help="Run in WebSocket mode")
    parser.add_argument("--port", type=int, default=3001, help="WebSocket port (default 3001)")
    parser.add_argument(
        "--token", type=str, default=None, help="Auth token (auto-generated if omitted)"
    )
    args = parser.parse_args()

    mode = "websocket" if args.websocket or args.port != 3001 else "stdio"
    start_mcp_server(mode=mode, port=args.port, token=args.token)
