"""
Brynq Knowledge MCP Server — Codebase intelligence for AI agents.

Lets any MCP-compatible agent (Claude Code, Gemini CLI, Cursor, etc.)
query the Brynq codebase: search code, find symbols, read modules,
list architecture, browse docs, and discover API endpoints.

Usage:
  Add to .mcp.json:
  {
    "brynq-knowledge": {
      "command": "python",
      "args": ["-m", "brynq.mcp_servers.knowledge_mcp"]
    }
  }
"""

import ast
import asyncio
import json
import os
import re
import subprocess
from pathlib import Path
from typing import Any, Dict, List, Optional

# Project root — 3 levels up from this file
_PROJECT_ROOT = Path(__file__).resolve().parents[3]


class SymbolIndex:
    """Pre-built index of classes, functions, and their locations."""

    def __init__(self, root: Path):
        self.root = root
        self._symbols: List[Dict[str, str]] = []
        self._built = False

    def build(self) -> None:
        """Walk Python files and extract class/function definitions."""
        self._symbols = []
        for py_dir in ["runtime", "cloud", "src/brynq"]:
            base = self.root / py_dir
            if not base.exists():
                continue
            for py_file in base.rglob("*.py"):
                if "__pycache__" in str(py_file) or "_future" in str(py_file) or "_shelved" in str(py_file):
                    continue
                rel = str(py_file.relative_to(self.root)).replace("\\", "/")
                try:
                    source = py_file.read_text("utf-8", errors="replace")
                    tree = ast.parse(source, filename=str(py_file))
                    for node in ast.walk(tree):
                        if isinstance(node, ast.ClassDef):
                            self._symbols.append({
                                "type": "class",
                                "name": node.name,
                                "file": rel,
                                "line": node.lineno,
                            })
                        elif isinstance(node, ast.FunctionDef) or isinstance(node, ast.AsyncFunctionDef):
                            # Skip private/dunder methods unless they're top-level
                            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                                self._symbols.append({
                                    "type": "function",
                                    "name": node.name,
                                    "file": rel,
                                    "line": node.lineno,
                                })
                except (SyntaxError, UnicodeDecodeError):
                    continue
        self._built = True

    def search(self, query: str, limit: int = 20) -> List[Dict[str, str]]:
        """Search symbols by name (substring match, case-insensitive)."""
        if not self._built:
            self.build()
        q = query.lower()
        matches = [s for s in self._symbols if q in s["name"].lower()]
        return matches[:limit]

    def list_all(self, symbol_type: Optional[str] = None) -> List[Dict[str, str]]:
        """List all symbols, optionally filtered by type."""
        if not self._built:
            self.build()
        if symbol_type:
            return [s for s in self._symbols if s["type"] == symbol_type]
        return self._symbols


class BrynqKnowledge:
    """Core knowledge engine — the brains behind the MCP tools."""

    def __init__(self, root: Optional[Path] = None):
        self.root = root or _PROJECT_ROOT
        self.index = SymbolIndex(self.root)

    def search_code(self, pattern: str, file_glob: str = "*.py", max_results: int = 30) -> str:
        """Grep across the codebase for a pattern."""
        results = []
        try:
            for search_dir in ["runtime", "cloud", "src/brynq"]:
                base = self.root / search_dir
                if not base.exists():
                    continue
                for py_file in base.rglob(file_glob):
                    if "__pycache__" in str(py_file) or "_future" in str(py_file) or "_shelved" in str(py_file):
                        continue
                    try:
                        lines = py_file.read_text("utf-8", errors="replace").splitlines()
                        for i, line in enumerate(lines, 1):
                            if re.search(pattern, line, re.IGNORECASE):
                                rel = str(py_file.relative_to(self.root)).replace("\\", "/")
                                results.append(f"{rel}:{i}: {line.strip()}")
                                if len(results) >= max_results:
                                    return "\n".join(results)
                    except Exception:
                        continue
        except Exception as e:
            return f"Error: {e}"
        return "\n".join(results) if results else "No matches found."

    def find_symbol(self, name: str, limit: int = 20) -> str:
        """Find a class or function by name."""
        matches = self.index.search(name, limit)
        if not matches:
            return f"No symbol matching '{name}' found."
        lines = []
        for m in matches:
            lines.append(f"{m['type']:10s} {m['name']:40s} {m['file']}:{m['line']}")
        return "\n".join(lines)

    def read_module(self, file_path: str, start_line: int = 1, end_line: int = 200) -> str:
        """Read a specific file from the codebase."""
        full_path = self.root / file_path
        if not full_path.exists():
            return f"File not found: {file_path}"
        # Security: must be within project root
        try:
            full_path.resolve().relative_to(self.root.resolve())
        except ValueError:
            return "Access denied: path outside project root."
        try:
            lines = full_path.read_text("utf-8", errors="replace").splitlines()
            selected = lines[start_line - 1:end_line]
            numbered = [f"{i + start_line:4d} | {line}" for i, line in enumerate(selected)]
            return "\n".join(numbered)
        except Exception as e:
            return f"Error reading file: {e}"

    def list_modules(self, package: str = "runtime") -> str:
        """List all modules in a package with their main classes."""
        base = self.root / package
        if not base.exists():
            return f"Package '{package}' not found."
        result = []
        for py_file in sorted(base.rglob("*.py")):
            if "__pycache__" in str(py_file) or py_file.name == "__init__.py":
                continue
            if "_future" in str(py_file) or "_shelved" in str(py_file):
                continue
            rel = str(py_file.relative_to(self.root)).replace("\\", "/")
            # Extract class names
            try:
                source = py_file.read_text("utf-8", errors="replace")
                tree = ast.parse(source, filename=str(py_file))
                classes = [n.name for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]
                class_str = ", ".join(classes) if classes else "(functions only)"
                result.append(f"  {rel:55s} {class_str}")
            except (SyntaxError, UnicodeDecodeError):
                result.append(f"  {rel:55s} (parse error)")
        return f"{package}/ — {len(result)} modules:\n" + "\n".join(result)

    def get_architecture(self) -> str:
        """Return the CLAUDE.md architecture overview."""
        claude_md = self.root / "CLAUDE.md"
        if not claude_md.exists():
            return "CLAUDE.md not found."
        return claude_md.read_text("utf-8", errors="replace")

    def get_docs(self, query: str) -> str:
        """Search docs/ directory for a topic."""
        docs_dir = self.root / "docs"
        if not docs_dir.exists():
            return "docs/ directory not found."
        results = []
        q = query.lower()
        for md_file in docs_dir.rglob("*.md"):
            try:
                content = md_file.read_text("utf-8", errors="replace")
                if q in content.lower():
                    rel = str(md_file.relative_to(self.root)).replace("\\", "/")
                    # Find matching lines
                    for i, line in enumerate(content.splitlines(), 1):
                        if q in line.lower():
                            results.append(f"{rel}:{i}: {line.strip()}")
                            if len(results) >= 20:
                                return "\n".join(results)
            except Exception:
                continue
        return "\n".join(results) if results else f"No docs matching '{query}'."

    def get_endpoints(self) -> str:
        """List all cloud API endpoints from cloud/server.py."""
        server_py = self.root / "cloud" / "server.py"
        if not server_py.exists():
            return "cloud/server.py not found."
        try:
            source = server_py.read_text("utf-8", errors="replace")
            endpoints = []
            for i, line in enumerate(source.splitlines(), 1):
                match = re.match(r'\s+@app\.(get|post|put|delete|patch|websocket)\(["\'](.+?)["\']', line)
                if match:
                    method = match.group(1).upper()
                    path = match.group(2)
                    endpoints.append(f"  {method:10s} {path}")
            return f"Cloud API — {len(endpoints)} endpoints:\n" + "\n".join(endpoints)
        except Exception as e:
            return f"Error: {e}"

    def semantic_search(self, query: str, top_k: int = 5) -> str:
        """Semantic search using BGE embeddings. Requires built knowledge base."""
        kb_dir = self.root / "data" / "knowledge_base"
        index_path = kb_dir / "vector_index.npz"
        metadata_path = kb_dir / "metadata.jsonl"
        if not index_path.exists() or not metadata_path.exists():
            return "Knowledge base not built. Run: python scripts/build_knowledge_base.py"
        try:
            import numpy as np
            vectors = np.load(str(index_path))["vectors"]
            metadata = []
            for line in metadata_path.read_text("utf-8").strip().splitlines():
                metadata.append(json.loads(line))
            query_vec = self._embed_query(query)
            if query_vec is None:
                return "Embedding model not available. Install: pip install torch transformers"
            query_norm = query_vec / (np.linalg.norm(query_vec) + 1e-8)
            norms = np.linalg.norm(vectors, axis=1, keepdims=True) + 1e-8
            normalized = vectors / norms
            scores = normalized @ query_norm
            top_indices = np.argsort(scores)[::-1][:top_k]
            results = []
            for idx in top_indices:
                score = float(scores[idx])
                if score < 0.3:
                    continue
                meta = metadata[idx]
                source = meta["source"]
                offset = meta.get("char_offset", 0)
                fp = self.root / source
                preview = ""
                if fp.exists():
                    try:
                        text = fp.read_text("utf-8", errors="replace")
                        preview = text[offset:offset + 500].strip()
                    except Exception:
                        preview = "(could not read)"
                results.append("[" + str(round(score, 3)) + "] " + source + "\n" + preview + "\n")
            return "\n---\n".join(results) if results else "No relevant results found."
        except Exception as e:
            return "Semantic search error: " + str(e)

    def _embed_query(self, query: str):
        """Embed a query using BGE model. Returns numpy vector or None."""
        try:
            import torch
            from transformers import AutoModel, AutoTokenizer
            import numpy as np
            if not hasattr(self, '_query_tokenizer') or self._query_tokenizer is None:
                model_name = "BAAI/bge-base-en-v1.5"
                self._query_tokenizer = AutoTokenizer.from_pretrained(model_name)
                self._query_model = AutoModel.from_pretrained(model_name)
                device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
                self._query_model = self._query_model.to(device)
            tokenizer = self._query_tokenizer
            model = self._query_model
            device = next(model.parameters()).device
            prefixed = "Represent this sentence for searching relevant passages: " + query
            encoded = tokenizer([prefixed], padding=True, truncation=True, max_length=512, return_tensors="pt")
            encoded = {k: v.to(device) for k, v in encoded.items()}
            with torch.no_grad():
                outputs = model(**encoded)
            attention_mask = encoded["attention_mask"]
            token_embeddings = outputs.last_hidden_state
            mask_expanded = attention_mask.unsqueeze(-1).float()
            sum_embeddings = (token_embeddings * mask_expanded).sum(dim=1)
            count = mask_expanded.sum(dim=1).clamp(min=1e-9)
            mean_pooled = sum_embeddings / count
            norms = mean_pooled.norm(dim=1, keepdim=True).clamp(min=1e-8)
            normalized = mean_pooled / norms
            return normalized.squeeze(0).cpu().numpy().astype(np.float32)
        except ImportError:
            return None

    def get_file_tree(self, directory: str = "runtime", max_depth: int = 2) -> str:
        """Show the file tree of a directory."""
        base = self.root / directory
        if not base.exists():
            return f"Directory '{directory}' not found."
        lines = [f"{directory}/"]
        for item in sorted(base.rglob("*")):
            if "__pycache__" in str(item) or "_future" in str(item) or "_shelved" in str(item):
                continue
            rel = item.relative_to(base)
            depth = len(rel.parts)
            if depth > max_depth:
                continue
            indent = "  " * depth
            name = rel.name
            if item.is_dir():
                lines.append(f"{indent}{name}/")
            elif name.endswith(".py"):
                lines.append(f"{indent}{name}")
        return "\n".join(lines)

    def search_conversations(self, query: str, limit: int = 10) -> str:
        """Search past Brynq chat conversations by keyword.

        Returns matching sessions with snippets using FTS5 full-text search.
        """
        if not query or not query.strip():
            return "No query provided."
        try:
            from runtime.data.chat_db import ChatDB
            db = ChatDB()
            sessions = db.search_sessions(query, limit=limit)
            db.close()
        except Exception as e:
            return f"Search error: {e}"
        if not sessions:
            return f"No conversations matching '{query}'."
        lines = [f"Found {len(sessions)} session(s) matching '{query}':\n"]
        for s in sessions:
            title = s.get("session_title", "Untitled")
            date = s.get("updated_at", "")[:10]
            count = s.get("match_count", 0)
            snippet = s.get("best_snippet", "")
            lines.append(f'  Session: "{title}" ({date})  [{count} match{"es" if count != 1 else ""}]')
            if snippet:
                lines.append(f"  ...{snippet}...")
            lines.append("")
        return "\n".join(lines)


# ── MCP Server (stdio protocol) ────────────────────────────────────

def _build_tool_list() -> List[Dict[str, Any]]:
    """Return MCP tool definitions."""
    return [
        {
            "name": "search_code",
            "description": "Search the Brynq codebase for a regex pattern. Returns matching lines with file paths and line numbers.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "pattern": {"type": "string", "description": "Regex pattern to search for"},
                    "file_glob": {"type": "string", "description": "File glob pattern (default: *.py)", "default": "*.py"},
                    "max_results": {"type": "integer", "description": "Maximum results to return", "default": 30},
                },
                "required": ["pattern"],
            },
        },
        {
            "name": "find_symbol",
            "description": "Find a class or function definition by name. Returns file path and line number.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "name": {"type": "string", "description": "Class or function name to find (substring match)"},
                },
                "required": ["name"],
            },
        },
        {
            "name": "read_module",
            "description": "Read a specific file from the Brynq codebase. Returns numbered lines.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "file_path": {"type": "string", "description": "Relative path from project root (e.g., runtime/rooms/room_host.py)"},
                    "start_line": {"type": "integer", "description": "Start line (default: 1)", "default": 1},
                    "end_line": {"type": "integer", "description": "End line (default: 200)", "default": 200},
                },
                "required": ["file_path"],
            },
        },
        {
            "name": "list_modules",
            "description": "List all Python modules in a package (runtime, cloud, or src/brynq) with their main classes.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "package": {"type": "string", "description": "Package to list: runtime, cloud, or src/brynq", "default": "runtime"},
                },
            },
        },
        {
            "name": "get_architecture",
            "description": "Get the full Brynq architecture overview from CLAUDE.md. Includes project structure, data flow rules, stack info.",
            "inputSchema": {"type": "object", "properties": {}},
        },
        {
            "name": "get_docs",
            "description": "Search the docs/ directory for a topic. Returns matching lines from markdown files.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Topic to search for in documentation"},
                },
                "required": ["query"],
            },
        },
        {
            "name": "get_endpoints",
            "description": "List all cloud API endpoints from cloud/server.py with HTTP methods and paths.",
            "inputSchema": {"type": "object", "properties": {}},
        },
        {
            "name": "semantic_search",
            "description": "Semantic search across the entire Brynq codebase using AI embeddings. Understands meaning, not just keywords. Example: 'how does authentication work?' finds auth code even without exact keyword matches. Requires: python scripts/build_knowledge_base.py to be run first.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Natural language search query"},
                    "top_k": {"type": "integer", "description": "Number of results (default: 5)", "default": 5},
                },
                "required": ["query"],
            },
        },
        {
            "name": "get_file_tree",
            "description": "Show the directory tree of a package. Good for understanding project structure.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "directory": {"type": "string", "description": "Directory to show (e.g., runtime, cloud, app/src)", "default": "runtime"},
                    "max_depth": {"type": "integer", "description": "Maximum depth to show", "default": 2},
                },
            },
        },
        {
            "name": "search_conversations",
            "description": "Search past Brynq chat conversations by keyword. Returns matching sessions with snippets.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Search query"},
                    "limit": {"type": "integer", "description": "Maximum sessions to return (default: 10)", "default": 10},
                },
                "required": ["query"],
            },
        },
    ]


async def _handle_request(knowledge: BrynqKnowledge, method: str, params: Dict[str, Any]) -> Any:
    """Route MCP requests to knowledge engine methods."""
    if method == "initialize":
        return {
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {"listChanged": False}},
            "serverInfo": {"name": "brynq-knowledge", "version": "0.420.69"},
        }
    elif method == "tools/list":
        return {"tools": _build_tool_list()}
    elif method == "tools/call":
        tool_name = params.get("name", "")
        args = params.get("arguments", {})

        if tool_name == "search_code":
            result = knowledge.search_code(args["pattern"], args.get("file_glob", "*.py"), args.get("max_results", 30))
        elif tool_name == "find_symbol":
            result = knowledge.find_symbol(args["name"])
        elif tool_name == "read_module":
            result = knowledge.read_module(args["file_path"], args.get("start_line", 1), args.get("end_line", 200))
        elif tool_name == "list_modules":
            result = knowledge.list_modules(args.get("package", "runtime"))
        elif tool_name == "get_architecture":
            result = knowledge.get_architecture()
        elif tool_name == "get_docs":
            result = knowledge.get_docs(args["query"])
        elif tool_name == "get_endpoints":
            result = knowledge.get_endpoints()
        elif tool_name == "semantic_search":
            result = knowledge.semantic_search(args["query"], args.get("top_k", 5))
        elif tool_name == "get_file_tree":
            result = knowledge.get_file_tree(args.get("directory", "runtime"), args.get("max_depth", 2))
        elif tool_name == "search_conversations":
            result = knowledge.search_conversations(args["query"], args.get("limit", 10))
        else:
            result = f"Unknown tool: {tool_name}"

        return {"content": [{"type": "text", "text": result}]}
    elif method == "notifications/initialized":
        return None
    else:
        return {"error": {"code": -32601, "message": f"Method not found: {method}"}}


async def _run_stdio():
    """Run the MCP server over stdio (JSON-RPC)."""
    import sys

    knowledge = BrynqKnowledge()
    # Pre-build symbol index on startup
    knowledge.index.build()

    reader = asyncio.StreamReader()
    protocol = asyncio.StreamReaderProtocol(reader)
    await asyncio.get_event_loop().connect_read_pipe(lambda: protocol, sys.stdin.buffer)

    writer_transport, writer_protocol = await asyncio.get_event_loop().connect_write_pipe(
        asyncio.streams.FlowControlMixin, sys.stdout.buffer
    )
    writer = asyncio.StreamWriter(writer_transport, writer_protocol, None, asyncio.get_event_loop())

    while True:
        # Read Content-Length header
        header = b""
        while True:
            byte = await reader.read(1)
            if not byte:
                return  # EOF
            header += byte
            if header.endswith(b"\r\n\r\n"):
                break

        # Parse content length
        match = re.search(rb"Content-Length: (\d+)", header)
        if not match:
            continue
        length = int(match.group(1))

        # Read body
        body = await reader.readexactly(length)
        request = json.loads(body.decode("utf-8"))

        msg_id = request.get("id")
        method = request.get("method", "")
        params = request.get("params", {})

        result = await _handle_request(knowledge, method, params)

        if result is None:
            continue  # Notification, no response needed

        response = {"jsonrpc": "2.0", "id": msg_id, "result": result}
        response_bytes = json.dumps(response).encode("utf-8")
        header_bytes = f"Content-Length: {len(response_bytes)}\r\n\r\n".encode("utf-8")

        writer.write(header_bytes + response_bytes)
        await writer.drain()


def run():
    """Entry point for MCP server."""
    asyncio.run(_run_stdio())


if __name__ == "__main__":
    run()
