"""TokenNuke MCP Server — 13 tools for intelligent code exploration.

Nuke your token usage. Index codebases with tree-sitter, search with
hybrid FTS5+vector, trace call graphs, and retrieve symbols with O(1)
byte-offset reads.
"""

import hashlib
import json
import logging
import os
import re
from pathlib import Path
from typing import Optional

import pathspec
from mcp.server.fastmcp import FastMCP

from tokennuke.parser.extractor import parse_file
from tokennuke.parser.call_graph import extract_calls
from tokennuke.parser.languages import LANGUAGE_EXTENSIONS, LANGUAGE_REGISTRY
from tokennuke.parser.symbols import compute_file_hash
from tokennuke.security import should_exclude_file, is_binary_file, safe_decode
from tokennuke.storage.database import Database
from tokennuke.embedder.embed import Embedder, hybrid_search, get_embedder

log = logging.getLogger(__name__)

# Global state
_data_dir: Path | None = None
_repos: dict[str, dict] = {}  # repo_id -> {path, db, name}
_embedder: Embedder | None = None


def _get_data_dir() -> Path:
    """Get or create the data directory."""
    global _data_dir
    if _data_dir is None:
        _data_dir = Path.home() / '.tokennuke'
    _data_dir.mkdir(parents=True, exist_ok=True)
    return _data_dir


def _repo_id_from_path(path: str) -> str:
    """Generate a stable repo ID from a path."""
    return hashlib.md5(path.encode()).hexdigest()[:12]


def _get_db(repo_id: str) -> Database | None:
    """Get the database for a repo."""
    if repo_id in _repos:
        return _repos[repo_id]['db']
    # Check if DB file exists on disk
    db_path = _get_data_dir() / 'repos' / f'{repo_id}.db'
    if db_path.exists():
        db = Database(db_path)
        _repos[repo_id] = {'db': db, 'path': '', 'name': repo_id}
        return db
    return None


def _get_embedder_safe() -> Embedder | None:
    """Get embedder, returning None if fastembed is not installed."""
    global _embedder
    if _embedder is not None:
        return _embedder
    try:
        _embedder = get_embedder()
        return _embedder
    except (ImportError, Exception):
        return None


def _walk_source_files(
    root: Path,
    include_patterns: list[str] | None = None,
    exclude_patterns: list[str] | None = None,
) -> list[Path]:
    """Walk a directory and return source files to index.

    Respects .gitignore if present. Applies include/exclude glob filters.
    """
    root = root.resolve()
    files: list[Path] = []

    # Load .gitignore
    gitignore_spec = None
    gitignore_path = root / '.gitignore'
    if gitignore_path.exists():
        try:
            with open(gitignore_path) as f:
                gitignore_spec = pathspec.PathSpec.from_lines('gitwildmatch', f)
        except Exception:
            pass

    # Always exclude these directories
    skip_dirs = {
        '.git', '.svn', '.hg', 'node_modules', '__pycache__', '.mypy_cache',
        '.pytest_cache', '.tox', '.venv', 'venv', 'env', '.env',
        'dist', 'build', '.next', '.nuxt', 'target', 'vendor',
        '.tokennuke', '.claude',
    }

    include_spec = None
    exclude_spec = None
    if include_patterns:
        include_spec = pathspec.PathSpec.from_lines('gitwildmatch', include_patterns)
    if exclude_patterns:
        exclude_spec = pathspec.PathSpec.from_lines('gitwildmatch', exclude_patterns)

    for dirpath, dirnames, filenames in os.walk(root):
        # Prune skipped directories
        dirnames[:] = [
            d for d in dirnames
            if d not in skip_dirs and not d.startswith('.')
        ]

        for filename in filenames:
            filepath = Path(dirpath) / filename
            rel_path = filepath.relative_to(root).as_posix()

            # Check gitignore
            if gitignore_spec and gitignore_spec.match_file(rel_path):
                continue

            # Check include filter
            if include_spec and not include_spec.match_file(rel_path):
                continue

            # Check exclude filter
            if exclude_spec and exclude_spec.match_file(rel_path):
                continue

            # Check file extension
            ext = filepath.suffix.lower()
            if ext not in LANGUAGE_EXTENSIONS:
                continue

            # Security checks
            reason = should_exclude_file(filepath, root)
            if reason:
                continue

            files.append(filepath)

    return files


def _index_files(
    db: Database,
    root: Path,
    files: list[Path],
    embed: bool = True,
) -> dict:
    """Index a list of source files into the database.

    Uses incremental indexing — skips files whose SHA-256 hasn't changed.
    Returns stats about what was indexed.
    """
    existing_hashes = db.get_file_hashes()
    root_resolved = root.resolve()

    stats = {
        'indexed': 0,
        'skipped': 0,
        'errors': 0,
        'deleted': 0,
        'symbols': 0,
        'call_edges': 0,
    }

    current_paths: set[str] = set()
    embedder = _get_embedder_safe() if embed else None
    embed_queue: list[tuple[int, str]] = []  # (symbol_rowid, text)

    for filepath in files:
        try:
            rel_path = filepath.relative_to(root_resolved).as_posix()
            current_paths.add(rel_path)

            # Incremental: skip unchanged files
            file_hash = compute_file_hash(str(filepath))
            if existing_hashes.get(rel_path) == file_hash:
                stats['skipped'] += 1
                continue

            # Read and parse
            with open(filepath, 'rb') as f:
                raw = f.read()

            content = safe_decode(raw)
            ext = filepath.suffix.lower()
            language = LANGUAGE_EXTENSIONS.get(ext)
            if not language:
                continue

            # Extract symbols
            symbols = parse_file(content, rel_path, language)

            # Extract call graph edges per function/method
            spec = LANGUAGE_REGISTRY.get(language)
            call_edges: dict[str, list] = {}
            if spec:
                source_bytes = content.encode('utf-8')
                from tree_sitter_language_pack import get_parser
                parser = get_parser(spec.ts_language)
                tree = parser.parse(source_bytes)

                for sym in symbols:
                    if sym.kind in ('function', 'method'):
                        # Find the AST node for this symbol by byte offset
                        body_bytes = source_bytes[sym.byte_offset:sym.byte_offset + sym.byte_length]
                        # Re-parse just this symbol's region isn't feasible,
                        # so walk the full tree to find the matching node
                        sym_node = _find_node_at(
                            tree.root_node, sym.byte_offset, sym.byte_offset + sym.byte_length,
                        )
                        if sym_node:
                            edges = extract_calls(sym_node, spec, source_bytes)
                            if edges:
                                call_edges[sym.id] = edges
                                stats['call_edges'] += len(edges)

            # Store
            file_id = db.upsert_file(
                path=rel_path,
                sha256=file_hash,
                language=language,
                size_bytes=filepath.stat().st_size,
                content=content,
                symbols=symbols,
                call_edges=call_edges,
            )

            stats['indexed'] += 1
            stats['symbols'] += len(symbols)

            # Queue embeddings
            if embedder:
                for sym in symbols:
                    text = embedder.embed_symbol({
                        'kind': sym.kind,
                        'qualified_name': sym.qualified_name,
                        'signature': sym.signature,
                        'docstring': sym.docstring,
                    })
                    # We need the rowid from the symbols table
                    row = db.conn.execute(
                        'SELECT id FROM symbols WHERE symbol_id = ?',
                        (sym.id,),
                    ).fetchone()
                    if row:
                        embed_queue.append((row['id'], text))

        except Exception:
            log.warning('Error indexing %s', filepath, exc_info=True)
            stats['errors'] += 1

    # Delete removed files
    for old_path in existing_hashes:
        if old_path not in current_paths:
            db.delete_file(old_path)
            stats['deleted'] += 1

    # Resolve call graph edges
    resolved = db.resolve_call_edges()
    stats['call_edges_resolved'] = resolved

    # Batch embed
    if embed_queue and embedder:
        try:
            texts = [t for _, t in embed_queue]
            embeddings = embedder.embed_batch(texts)
            pairs = [(rowid, emb) for (rowid, _), emb in zip(embed_queue, embeddings)]
            db.upsert_embeddings_batch(pairs)
            stats['embeddings'] = len(pairs)
        except Exception:
            log.warning('Embedding failed', exc_info=True)
            stats['embeddings'] = 0

    return stats


def _find_node_at(root_node, start_byte: int, end_byte: int):
    """Find the AST node that exactly matches the given byte range."""
    if root_node.start_byte == start_byte and root_node.end_byte == end_byte:
        return root_node
    for child in root_node.children:
        if child.start_byte <= start_byte and child.end_byte >= end_byte:
            result = _find_node_at(child, start_byte, end_byte)
            if result:
                return result
    return None


def create_server(data_dir: str | None = None) -> FastMCP:
    """Create and configure the TokenNuke MCP server."""
    global _data_dir

    if data_dir:
        _data_dir = Path(data_dir)

    mcp = FastMCP(
        'TokenNuke',
        instructions=(
            'TokenNuke indexes codebases and provides O(1) symbol retrieval. '
            'Start by calling index_folder to index a codebase, then use '
            'search_symbols to find code, get_symbol to read it, and '
            'get_callers/get_callees to trace the call graph.'
        ),
    )

    # --- Tool 1: index_folder ---
    @mcp.tool()
    def index_folder(
        path: str,
        include_patterns: list[str] | None = None,
        exclude_patterns: list[str] | None = None,
        embed: bool = True,
    ) -> str:
        """Index a local directory's source code for fast symbol retrieval.

        Scans the directory for source files (Python, JS, TS, Go, Rust, Java,
        C, C++, C#, Ruby), extracts symbols with byte offsets, builds a call
        graph, and creates search indexes. Uses incremental indexing — only
        re-indexes files that changed since last run.

        Args:
            path: Absolute path to the directory to index.
            include_patterns: Optional glob patterns to include (e.g., ["src/**"]).
            exclude_patterns: Optional glob patterns to exclude (e.g., ["*_test.go"]).
            embed: Whether to compute embeddings for semantic search (slower but better search).

        Returns:
            JSON with repo_id and indexing stats.
        """
        folder = Path(path).resolve()
        if not folder.is_dir():
            return json.dumps({'error': f'Not a directory: {path}'})

        repo_id = _repo_id_from_path(str(folder))
        db_path = _get_data_dir() / 'repos' / f'{repo_id}.db'
        db = Database(db_path)

        # Store repo metadata
        db.conn.execute(
            'INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)',
            ('repo_path', str(folder)),
        )
        db.conn.execute(
            'INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)',
            ('repo_name', folder.name),
        )
        db.conn.commit()

        _repos[repo_id] = {'db': db, 'path': str(folder), 'name': folder.name}

        # Walk and index
        files = _walk_source_files(folder, include_patterns, exclude_patterns)
        stats = _index_files(db, folder, files, embed=embed)
        stats['repo_id'] = repo_id
        stats['repo_name'] = folder.name
        stats['total_files'] = len(files)
        stats.update(db.get_stats())

        return json.dumps(stats, indent=2)

    # --- Tool 2: index_repo ---
    @mcp.tool()
    def index_repo(
        url: str,
        branch: str = 'main',
        sparse_paths: list[str] | None = None,
    ) -> str:
        """Index a remote Git repository (GitHub, GitLab, Bitbucket).

        Clones the repo (or sparse-checks out specific paths) into a
        temporary directory and indexes it. Supports any Git URL.

        Args:
            url: Git clone URL (https or ssh).
            branch: Branch to index (default: main).
            sparse_paths: Optional paths to sparse-checkout (saves bandwidth).

        Returns:
            JSON with repo_id and indexing stats.
        """
        import subprocess
        import tempfile

        # Sanitize URL
        if not re.match(r'^https?://|^git@', url):
            return json.dumps({'error': 'Invalid git URL'})

        clone_dir = Path(tempfile.mkdtemp(prefix='tokennuke-'))
        try:
            cmd = ['git', 'clone', '--depth', '1', '--branch', branch]

            if sparse_paths:
                cmd.extend(['--sparse', '--filter=blob:none'])

            cmd.extend([url, str(clone_dir / 'repo')])
            subprocess.run(cmd, check=True, capture_output=True, timeout=120)

            repo_dir = clone_dir / 'repo'

            if sparse_paths:
                subprocess.run(
                    ['git', 'sparse-checkout', 'set'] + sparse_paths,
                    cwd=str(repo_dir),
                    check=True,
                    capture_output=True,
                )

            # Index the cloned repo
            return index_folder(str(repo_dir))

        except subprocess.TimeoutExpired:
            return json.dumps({'error': 'Clone timed out (120s limit)'})
        except subprocess.CalledProcessError as e:
            return json.dumps({'error': f'Git clone failed: {e.stderr.decode()[:200]}'})
        finally:
            import shutil
            shutil.rmtree(clone_dir, ignore_errors=True)

    # --- Tool 3: list_repos ---
    @mcp.tool()
    def list_repos() -> str:
        """List all indexed repositories.

        Returns:
            JSON array of repos with id, name, path, and stats.
        """
        repos_dir = _get_data_dir() / 'repos'
        if not repos_dir.exists():
            return json.dumps([])

        result = []
        for db_file in repos_dir.glob('*.db'):
            repo_id = db_file.stem
            db = _get_db(repo_id)
            if db:
                try:
                    name_row = db.conn.execute(
                        "SELECT value FROM meta WHERE key='repo_name'",
                    ).fetchone()
                    path_row = db.conn.execute(
                        "SELECT value FROM meta WHERE key='repo_path'",
                    ).fetchone()
                    stats = db.get_stats()
                    result.append({
                        'repo_id': repo_id,
                        'name': name_row['value'] if name_row else repo_id,
                        'path': path_row['value'] if path_row else '',
                        **stats,
                    })
                except Exception:
                    result.append({'repo_id': repo_id, 'error': 'corrupt'})

        return json.dumps(result, indent=2)

    # --- Tool 4: invalidate_cache ---
    @mcp.tool()
    def invalidate_cache(repo_id: str) -> str:
        """Force re-index a repository by clearing its file hashes.

        Next call to index_folder will re-parse all files.

        Args:
            repo_id: Repository ID from list_repos or index_folder.

        Returns:
            Confirmation message.
        """
        db = _get_db(repo_id)
        if not db:
            return json.dumps({'error': f'Repo {repo_id} not found'})

        db.conn.execute('UPDATE files SET sha256 = ""')
        db.conn.commit()
        return json.dumps({'status': 'ok', 'message': 'Cache invalidated. Re-index to rebuild.'})

    # --- Tool 5: file_tree ---
    @mcp.tool()
    def file_tree(
        repo_id: str,
        path_prefix: str = '',
        depth: int | None = None,
    ) -> str:
        """Get directory tree of indexed files.

        Args:
            repo_id: Repository ID.
            path_prefix: Optional path prefix to filter (e.g., "src/").
            depth: Max directory depth to return.

        Returns:
            JSON array of file/directory entries.
        """
        db = _get_db(repo_id)
        if not db:
            return json.dumps({'error': f'Repo {repo_id} not found'})

        tree = db.file_tree(path_prefix, depth)
        return json.dumps(tree, indent=2)

    # --- Tool 6: file_outline ---
    @mcp.tool()
    def file_outline(repo_id: str, file_path: str) -> str:
        """List all symbols in a single file, ordered by line number.

        Args:
            repo_id: Repository ID.
            file_path: Relative file path within the repo.

        Returns:
            JSON array of symbols with name, kind, signature, line numbers.
        """
        db = _get_db(repo_id)
        if not db:
            return json.dumps({'error': f'Repo {repo_id} not found'})

        symbols = db.file_outline(file_path)
        return json.dumps([
            {
                'symbol_id': s['symbol_id'],
                'name': s['name'],
                'qualified_name': s['qualified_name'],
                'kind': s['kind'],
                'signature': s['signature'],
                'line': s['line'],
                'end_line': s['end_line'],
            }
            for s in symbols
        ], indent=2)

    # --- Tool 7: repo_outline ---
    @mcp.tool()
    def repo_outline(
        repo_id: str,
        kind_filter: str | None = None,
        limit: int = 500,
    ) -> str:
        """List all symbols in a repository (summary view).

        Args:
            repo_id: Repository ID.
            kind_filter: Optional filter by kind (function, class, method, type, constant, interface).
            limit: Max symbols to return (default 500).

        Returns:
            JSON array of symbol summaries.
        """
        db = _get_db(repo_id)
        if not db:
            return json.dumps({'error': f'Repo {repo_id} not found'})

        symbols = db.repo_outline(kind_filter, limit)
        return json.dumps(symbols, indent=2)

    # --- Tool 8: get_symbol ---
    @mcp.tool()
    def get_symbol(repo_id: str, qualified_name: str) -> str:
        """Get the full source code of a single symbol using O(1) byte-offset read.

        This is the key token-saving feature: instead of reading an entire file,
        we seek directly to the symbol's byte offset and read only its bytes.

        Args:
            repo_id: Repository ID.
            qualified_name: The symbol's qualified name (e.g., "MyClass.login") or
                           full symbol_id (e.g., "src/main.py::MyClass.login#method").

        Returns:
            JSON with symbol metadata and full source code.
        """
        db = _get_db(repo_id)
        if not db:
            return json.dumps({'error': f'Repo {repo_id} not found'})

        # Try as symbol_id first, then as qualified_name
        sym = db.get_symbol(qualified_name)
        if not sym:
            sym = db.get_symbol_by_qualified_name(qualified_name)
        if not sym:
            return json.dumps({'error': f'Symbol not found: {qualified_name}'})

        # O(1) byte-offset read
        repo_path_row = db.conn.execute(
            "SELECT value FROM meta WHERE key='repo_path'",
        ).fetchone()
        if not repo_path_row:
            return json.dumps({'error': 'Repo path not found in metadata'})

        source = ''
        file_path = Path(repo_path_row['value']) / sym['file_path']
        try:
            with open(file_path, 'rb') as f:
                f.seek(sym['byte_offset'])
                raw = f.read(sym['byte_length'])
                source = raw.decode('utf-8', errors='replace')
        except (OSError, KeyError):
            source = '[source file not accessible]'

        return json.dumps({
            'symbol_id': sym['symbol_id'],
            'name': sym['name'],
            'qualified_name': sym['qualified_name'],
            'kind': sym['kind'],
            'language': sym['language'],
            'file': sym['file_path'],
            'line': sym['line'],
            'end_line': sym['end_line'],
            'signature': sym['signature'],
            'docstring': sym.get('docstring', ''),
            'source': source,
        }, indent=2)

    # --- Tool 9: get_symbols ---
    @mcp.tool()
    def get_symbols(repo_id: str, qualified_names: list[str]) -> str:
        """Batch get multiple symbols' source code.

        More efficient than multiple get_symbol calls.

        Args:
            repo_id: Repository ID.
            qualified_names: List of qualified names or symbol IDs.

        Returns:
            JSON array of symbol objects with source code.
        """
        results = []
        for qn in qualified_names:
            result = json.loads(get_symbol(repo_id, qn))
            results.append(result)
        return json.dumps(results, indent=2)

    # --- Tool 10: search_symbols ---
    @mcp.tool()
    def search_symbols(
        repo_id: str,
        query: str,
        kind: str | None = None,
        limit: int = 20,
    ) -> str:
        """Hybrid search for symbols using FTS5 + vector similarity.

        Combines keyword matching (BM25) with semantic similarity (embeddings)
        via Reciprocal Rank Fusion for best results. Finds symbols even when
        the query doesn't exactly match the name.

        Args:
            repo_id: Repository ID.
            query: Search query (e.g., "authentication middleware", "parse config").
            kind: Optional filter by kind (function, class, method, type, constant, interface).
            limit: Max results (default 20).

        Returns:
            JSON array of matching symbols with relevance scores.
        """
        db = _get_db(repo_id)
        if not db:
            return json.dumps({'error': f'Repo {repo_id} not found'})

        embedder = _get_embedder_safe()
        results = hybrid_search(db, query, embedder=embedder, kind=kind, limit=limit)

        return json.dumps([
            {
                'symbol_id': r.get('symbol_id'),
                'name': r.get('name'),
                'qualified_name': r.get('qualified_name'),
                'kind': r.get('kind'),
                'file': r.get('file_path'),
                'line': r.get('line'),
                'signature': r.get('signature'),
                'score': r.get('score'),
            }
            for r in results
        ], indent=2)

    # --- Tool 11: search_text ---
    @mcp.tool()
    def search_text(
        repo_id: str,
        query: str,
        glob: str | None = None,
        limit: int = 20,
    ) -> str:
        """Full-text search in raw file contents.

        Searches through the actual source code text, not just symbol names.
        Useful for finding string literals, comments, config values, TODOs, etc.

        Args:
            repo_id: Repository ID.
            query: Text to search for.
            glob: Optional glob pattern to filter files (e.g., "*.py", "src/**/*.ts").
            limit: Max results (default 20).

        Returns:
            JSON array of matches with file path and surrounding context.
        """
        db = _get_db(repo_id)
        if not db:
            return json.dumps({'error': f'Repo {repo_id} not found'})

        results = db.search_text(query, glob_pattern=glob, limit=limit)
        return json.dumps(results, indent=2)

    # --- Tool 12: get_callees ---
    @mcp.tool()
    def get_callees(
        repo_id: str,
        qualified_name: str,
        depth: int = 1,
    ) -> str:
        """Get functions/methods called by a symbol (outgoing call graph).

        Traces the call graph to show what a function calls. Useful for
        understanding dependencies and side effects.

        Args:
            repo_id: Repository ID.
            qualified_name: Symbol's qualified name or symbol_id.
            depth: How many levels deep to trace (default 1, max 5).

        Returns:
            JSON array of called functions with file and line info.
        """
        db = _get_db(repo_id)
        if not db:
            return json.dumps({'error': f'Repo {repo_id} not found'})

        depth = min(depth, 5)

        # Resolve to symbol_id
        sym = db.get_symbol(qualified_name)
        if not sym:
            sym = db.get_symbol_by_qualified_name(qualified_name)
        if not sym:
            return json.dumps({'error': f'Symbol not found: {qualified_name}'})

        callees = db.get_callees(sym['symbol_id'], depth=depth)
        return json.dumps(callees, indent=2)

    # --- Tool 13: get_callers ---
    @mcp.tool()
    def get_callers(
        repo_id: str,
        qualified_name: str,
        depth: int = 1,
    ) -> str:
        """Get functions/methods that call a symbol (incoming call graph).

        Traces the call graph backwards to show what calls a function.
        Useful for impact analysis and understanding usage patterns.

        Args:
            repo_id: Repository ID.
            qualified_name: Symbol's qualified name or symbol_id.
            depth: How many levels deep to trace (default 1, max 5).

        Returns:
            JSON array of calling functions with file and line info.
        """
        db = _get_db(repo_id)
        if not db:
            return json.dumps({'error': f'Repo {repo_id} not found'})

        depth = min(depth, 5)

        sym = db.get_symbol(qualified_name)
        if not sym:
            sym = db.get_symbol_by_qualified_name(qualified_name)
        if not sym:
            return json.dumps({'error': f'Symbol not found: {qualified_name}'})

        callers = db.get_callers(sym['symbol_id'], depth=depth)
        return json.dumps(callers, indent=2)

    return mcp
