"""Unified symbols tool implementation.

This module provides the SymbolsTool for searching, indexing, and querying code symbols
using tree-sitter AST parsing. It can find function definitions, class declarations,
and other code structures with full context.
"""

import os
from typing import (
    Any,
    Dict,
    List,
    Unpack,
    Optional,
    Annotated,
    TypedDict,
    final,
    override,
)
from pathlib import Path

from pydantic import Field
from grep_ast.grep_ast import TreeContext
from mcp.server.fastmcp import Context as MCPContext

from hanzo_mcp.tools.filesystem.base import FilesystemBaseTool
from hanzo_mcp.tools.common.auto_timeout import auto_timeout

# Parameter types
Action = Annotated[
    str,
    Field(
        description="Action: search (default), ast, index, query, list",
        default="search",
    ),
]

Pattern = Annotated[
    Optional[str],
    Field(
        description="Pattern to search for in code",
        default=None,
    ),
]

SearchPath = Annotated[
    str,
    Field(
        description="Path to search/index (file or directory)",
        default=".",
    ),
]

SymbolType = Annotated[
    Optional[str],
    Field(
        description="Symbol type: function, class, method, variable",
        default=None,
    ),
]

IgnoreCase = Annotated[
    bool,
    Field(
        description="Ignore case when matching",
        default=False,
    ),
]

ShowContext = Annotated[
    bool,
    Field(
        description="Show AST context around matches",
        default=True,
    ),
]

Limit = Annotated[
    int,
    Field(
        description="Maximum results to return",
        default=50,
    ),
]


class SymbolsParams(TypedDict, total=False):
    """Parameters for symbols tool."""

    action: str
    pattern: Optional[str]
    path: str
    symbol_type: Optional[str]
    ignore_case: bool
    show_context: bool
    limit: int


@final
class SymbolsTool(FilesystemBaseTool):
    """Tool for code symbol operations using tree-sitter."""

    def __init__(self, permission_manager):
        """Initialize the symbols tool."""
        super().__init__(permission_manager)
        self._symbol_cache = {}  # Cache for indexed symbols

    @property
    @override
    def name(self) -> str:
        """Get the tool name."""
        return "symbols"

    @property
    @override
    def description(self) -> str:
        """Get the tool description."""
        return """Code symbols search with tree-sitter AST. Actions: search (default), ast, index, query, list.

Usage:
symbols "function_name"
symbols --action ast --pattern "TODO" --path ./src
symbols --action query --symbol-type function --path ./src
symbols --action index --path ./project
symbols --action list --path ./src --symbol-type class

Finds code structures (functions, classes, methods) with full context."""

    @override
    @auto_timeout("symbols")
    async def call(
        self,
        ctx: MCPContext,
        **params: Unpack[SymbolsParams],
    ) -> str:
        """Execute symbols operation."""
        tool_ctx = self.create_tool_context(ctx)
        await self.set_tool_context_info(tool_ctx)

        # Extract action
        action = params.get("action", "search")

        # Route to appropriate handler
        if action == "search":
            return await self._handle_search(params, tool_ctx)
        elif action == "ast" or action == "grep_ast":  # Support both for backward compatibility
            return await self._handle_ast(params, tool_ctx)
        elif action == "index":
            return await self._handle_index(params, tool_ctx)
        elif action == "query":
            return await self._handle_query(params, tool_ctx)
        elif action == "list":
            return await self._handle_list(params, tool_ctx)
        else:
            return f"Error: Unknown action '{action}'. Valid actions: search, ast, index, query, list"

    async def _handle_search(self, params: Dict[str, Any], tool_ctx) -> str:
        """Search for pattern in code with AST context."""
        pattern = params.get("pattern")
        if not pattern:
            return "Error: pattern required for search action"

        path = params.get("path", ".")
        ignore_case = params.get("ignore_case", False)
        show_context = params.get("show_context", True)
        limit = params.get("limit", 50)

        # Validate path
        path_validation = self.validate_path(path)
        if not path_validation.is_valid:
            await tool_ctx.error(f"Invalid path: {path_validation.error_message}")
            return f"Error: Invalid path: {path_validation.error_message}"

        # Check permissions
        is_allowed, error_message = await self.check_path_allowed(path, tool_ctx)
        if not is_allowed:
            return error_message

        # Check existence
        is_exists, error_message = await self.check_path_exists(path, tool_ctx)
        if not is_exists:
            return error_message

        await tool_ctx.info(f"Searching for '{pattern}' in {path}")

        # Get files to process
        files_to_process = self._get_source_files(path)
        if not files_to_process:
            return f"No source code files found in {path}"

        # Process files
        results = []
        match_count = 0

        for file_path in files_to_process:
            if match_count >= limit:
                break

            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    code = f.read()

                tc = TreeContext(
                    file_path,
                    code,
                    color=False,
                    verbose=False,
                    line_number=True,
                )

                # Find matches
                loi = tc.grep(pattern, ignore_case)

                if loi:
                    if show_context:
                        tc.add_lines_of_interest(loi)
                        tc.add_context()
                        output = tc.format()
                    else:
                        # Just show matching lines
                        output = "\n".join([f"{line}: {code.splitlines()[line - 1]}" for line in loi])

                    results.append(f"\n{file_path}:\n{output}\n")
                    match_count += len(loi)

            except Exception as e:
                await tool_ctx.warning(f"Could not parse {file_path}: {str(e)}")

        if not results:
            return f"No matches found for '{pattern}' in {path}"

        output = [f"=== Symbol Search Results for '{pattern}' ==="]
        output.append(f"Found {match_count} matches in {len(results)} files\n")
        output.extend(results)

        if match_count >= limit:
            output.append(f"\n(Results limited to {limit} matches)")

        return "\n".join(output)

    async def _handle_ast(self, params: Dict[str, Any], tool_ctx) -> str:
        """AST-aware grep - shows code structure context around matches."""
        pattern = params.get("pattern")
        if not pattern:
            return "Error: pattern required for ast action"

        path = params.get("path", ".")
        ignore_case = params.get("ignore_case", False)
        show_context = params.get("show_context", True)
        limit = params.get("limit", 50)

        # Validate path
        path_validation = self.validate_path(path)
        if not path_validation.is_valid:
            await tool_ctx.error(f"Invalid path: {path_validation.error_message}")
            return f"Error: Invalid path: {path_validation.error_message}"

        # Check permissions
        is_allowed, error_message = await self.check_path_allowed(path, tool_ctx)
        if not is_allowed:
            return error_message

        # Check existence
        is_exists, error_message = await self.check_path_exists(path, tool_ctx)
        if not is_exists:
            return error_message

        await tool_ctx.info(f"Running AST-aware grep for '{pattern}' in {path}")

        # Get files to process
        files_to_process = self._get_source_files(path)
        if not files_to_process:
            return f"No source code files found in {path}"

        # Process files
        results = []
        match_count = 0

        for file_path in files_to_process:
            if match_count >= limit:
                break

            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    code = f.read()

                # Create TreeContext for AST parsing
                tc = TreeContext(
                    file_path,
                    code,
                    color=False,
                    verbose=False,
                    line_number=True,
                )

                # Find matches with case sensitivity option
                if ignore_case:
                    loi = tc.grep(pattern, ignore_case=True)
                else:
                    loi = tc.grep(pattern, ignore_case=False)

                if loi:
                    # Always show AST context for grep_ast
                    tc.add_lines_of_interest(loi)
                    tc.add_context()

                    # Get the formatted output with structure
                    output = tc.format()

                    # Add section separator and file info
                    results.append(f"\n{'=' * 60}")
                    results.append(f"File: {file_path}")
                    results.append(f"Matches: {len(loi)}")
                    results.append(f"{'=' * 60}\n")
                    results.append(output)

                    match_count += len(loi)

            except Exception as e:
                await tool_ctx.warning(f"Could not parse {file_path}: {str(e)}")

        if not results:
            return f"No matches found for '{pattern}' in {path}"

        output = [f"=== AST-aware Grep Results for '{pattern}' ==="]
        output.append(f"Total matches: {match_count} in {len([r for r in results if '===' in str(r)]) // 4} files\n")
        output.extend(results)

        if match_count >= limit:
            output.append(f"\n(Results limited to {limit} matches)")

        return "\n".join(output)

    async def _handle_index(self, params: Dict[str, Any], tool_ctx) -> str:
        """Index symbols in a codebase."""
        path = params.get("path", ".")

        # Validate path
        is_allowed, error_message = await self.check_path_allowed(path, tool_ctx)
        if not is_allowed:
            return error_message

        await tool_ctx.info(f"Indexing symbols in {path}...")

        files_to_process = self._get_source_files(path)
        if not files_to_process:
            return f"No source code files found in {path}"

        # Clear cache for this path
        self._symbol_cache[path] = {
            "functions": [],
            "classes": [],
            "methods": [],
            "variables": [],
        }

        indexed_count = 0
        symbol_count = 0

        for file_path in files_to_process:
            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    code = f.read()

                tc = TreeContext(file_path, code, color=False, verbose=False)

                # Extract symbols (simplified - would need proper tree-sitter queries)
                # This is a placeholder for actual symbol extraction
                symbols = self._extract_symbols(tc, file_path)

                for symbol_type, syms in symbols.items():
                    self._symbol_cache[path][symbol_type].extend(syms)
                    symbol_count += len(syms)

                indexed_count += 1

            except Exception as e:
                await tool_ctx.warning(f"Could not index {file_path}: {str(e)}")

        output = [f"=== Symbol Indexing Complete ==="]
        output.append(f"Indexed {indexed_count} files")
        output.append(f"Found {symbol_count} total symbols:")

        for symbol_type, symbols in self._symbol_cache[path].items():
            if symbols:
                output.append(f"  {symbol_type}: {len(symbols)}")

        return "\n".join(output)

    async def _handle_query(self, params: Dict[str, Any], tool_ctx) -> str:
        """Query indexed symbols."""
        path = params.get("path", ".")
        symbol_type = params.get("symbol_type")
        pattern = params.get("pattern")
        limit = params.get("limit", 50)

        # Check if we have indexed this path
        if path not in self._symbol_cache:
            return f"No symbols indexed for {path}. Run 'symbols --action index --path {path}' first."

        symbols = self._symbol_cache[path]
        results = []

        # Filter by type if specified
        if symbol_type:
            if symbol_type in symbols:
                candidates = symbols[symbol_type]
            else:
                return f"Unknown symbol type: {symbol_type}. Valid types: {', '.join(symbols.keys())}"
        else:
            # Combine all symbol types
            candidates = []
            for syms in symbols.values():
                candidates.extend(syms)

        # Filter by pattern if specified
        if pattern:
            filtered = []
            for sym in candidates:
                if pattern.lower() in sym["name"].lower():
                    filtered.append(sym)
            candidates = filtered

        # Limit results
        candidates = candidates[:limit]

        if not candidates:
            return "No symbols found matching criteria"

        output = [f"=== Symbol Query Results ==="]
        output.append(f"Found {len(candidates)} symbols\n")

        for sym in candidates:
            output.append(f"{sym['type']}: {sym['name']}")
            output.append(f"  File: {sym['file']}:{sym['line']}")
            if sym.get("signature"):
                output.append(f"  Signature: {sym['signature']}")
            output.append("")

        return "\n".join(output)

    async def _handle_list(self, params: Dict[str, Any], tool_ctx) -> str:
        """List all symbols in a path."""
        # Similar to query but shows all symbols
        params["pattern"] = None
        return await self._handle_query(params, tool_ctx)

    def _get_source_files(self, path: str) -> List[str]:
        """Get all source code files in a path."""
        path_obj = Path(path)
        files_to_process = []

        # Common source file extensions
        extensions = {
            ".py",
            ".js",
            ".ts",
            ".jsx",
            ".tsx",
            ".java",
            ".cpp",
            ".c",
            ".h",
            ".hpp",
            ".cs",
            ".rb",
            ".go",
            ".rs",
            ".swift",
            ".kt",
            ".scala",
            ".php",
            ".lua",
            ".r",
            ".jl",
            ".ex",
            ".exs",
            ".clj",
            ".cljs",
        }

        if path_obj.is_file():
            if path_obj.suffix in extensions:
                files_to_process.append(str(path_obj))
        elif path_obj.is_dir():
            for root, _, files in os.walk(path_obj):
                for file in files:
                    file_path = Path(root) / file
                    if file_path.suffix in extensions and self.is_path_allowed(str(file_path)):
                        files_to_process.append(str(file_path))

        return files_to_process

    def _extract_symbols(self, tc: TreeContext, file_path: str) -> Dict[str, List[Dict[str, Any]]]:
        """Extract symbols from a TreeContext (placeholder implementation)."""
        # This would need proper tree-sitter queries to extract symbols
        # For now, return empty structure
        return {
            "functions": [],
            "classes": [],
            "methods": [],
            "variables": [],
        }

    def register(self, mcp_server) -> None:
        """Register this tool with the MCP server."""
        pass
