"""Search documentation tool for GPU programming corpora.

Provides semantic and keyword search over documentation for CuTeDSL, CUDA, etc.

Corpora are downloaded via `wafer corpus download <name>` and stored in ~/.cache/wafer/corpora/.
"""

import re
from pathlib import Path

from wafer_core.rollouts.dtypes import Tool, ToolCall, ToolFunction, ToolFunctionParameter, ToolResult

# Cache directory where wafer corpus download stores files
CACHE_DIR = Path.home() / ".cache" / "wafer" / "corpora"

# Available corpora (names match wafer corpus download)
AVAILABLE_CORPORA = ["cutlass", "cutedsl", "cuda", "hip", "amd"]

SEARCH_DOCS_TOOL = Tool(
    type="function",
    function=ToolFunction(
        name="search_docs",
        description="""Search GPU programming documentation for relevant information.

Use this tool to find documentation about:
- CUTLASS C++ (cute:: namespace, gemm tutorials, tensor cores, TMA, Blackwell)
- CuTeDSL Python API (@cute.kernel, @cute.jit, cute.arch functions)
- CUDA programming concepts
- GPU kernel optimization techniques
- Code examples and patterns

Available corpora:
- 'cutlass' - NVIDIA CUTLASS C++ docs + GitHub examples (gemm, hopper, blackwell)
- 'cutedsl' - CuTeDSL Python documentation
- 'cuda' - General CUDA programming docs
- 'hip' - AMD HIP programming docs
- 'amd' - AMD GPU kernel development (rocWMMA, CK, etc.)

Note: Corpora must be downloaded first with `wafer corpus download <name>`.
Returns relevant documentation snippets with file paths.""",
        parameters=ToolFunctionParameter(
            type="object",
            properties={
                "query": {
                    "type": "string",
                    "description": "Search query - describe what you're looking for",
                },
                "corpus": {
                    "type": "string",
                    "description": "Which docs to search: 'cutlass', 'cutedsl', 'cuda', 'hip', 'amd' (default: cutlass)",
                },
                "max_results": {
                    "type": "integer",
                    "description": "Maximum number of results to return (default: 5)",
                },
            },
        ),
        required=["query"],
    )
)


def _get_corpus_path(corpus_name: str) -> Path | None:
    """Get the path to a corpus in the cache directory.
    
    Corpora are stored at ~/.cache/wafer/corpora/<corpus_name>/
    """
    if corpus_name not in AVAILABLE_CORPORA:
        return None
    
    corpus_path = CACHE_DIR / corpus_name
    if corpus_path.exists():
        return corpus_path
    
    return None


def _search_files(corpus_path: Path, query: str, max_results: int = 5) -> list[dict]:
    """Simple keyword search through documentation files."""
    results = []
    query_terms = query.lower().split()
    
    # Search .md, .py, .cu, .hpp, and .h files (for CUTLASS examples)
    for pattern in ["**/*.md", "**/*.py", "**/*.cu", "**/*.hpp", "**/*.h", "**/*.cuh"]:
        for file_path in corpus_path.glob(pattern):
            if file_path.is_file():
                try:
                    content = file_path.read_text(encoding="utf-8", errors="ignore")
                    content_lower = content.lower()
                    
                    # Score based on term matches
                    score = sum(content_lower.count(term) for term in query_terms)
                    
                    if score > 0:
                        # Extract relevant snippets
                        snippets = _extract_snippets(content, query_terms)
                        results.append({
                            "file": str(file_path),  # Return absolute path so read tool can access it
                            "score": score,
                            "snippets": snippets[:3],  # Top 3 snippets
                        })
                except Exception:
                    continue
    
    # Sort by score and return top results
    results.sort(key=lambda x: x["score"], reverse=True)
    return results[:max_results]


def _extract_snippets(content: str, terms: list[str], context_lines: int = 5) -> list[str]:
    """Extract snippets containing search terms."""
    snippets = []
    lines = content.split("\n")
    
    for i, line in enumerate(lines):
        line_lower = line.lower()
        if any(term in line_lower for term in terms):
            # Get context around the match
            start = max(0, i - context_lines)
            end = min(len(lines), i + context_lines + 1)
            snippet = "\n".join(lines[start:end])
            
            # Skip very short snippets
            if len(snippet.strip()) > 50:
                snippets.append(snippet)
    
    return snippets


async def exec_search_docs(
    tool_call: ToolCall,
    corpus_override: str | None = None,
) -> ToolResult:
    """Execute search_docs tool.
    
    Args:
        tool_call: The tool call with query and optional corpus
        corpus_override: Override corpus path (for testing)
    """
    query = tool_call.args.get("query", "")
    corpus_name = tool_call.args.get("corpus", "cutlass")
    max_results = tool_call.args.get("max_results", 5)
    
    if not query:
        return ToolResult(
            tool_call_id=tool_call.id,
            content="",
            error="query parameter is required",
        )
    
    # Find corpus path
    if corpus_override:
        corpus_path = Path(corpus_override)
    else:
        corpus_path = _get_corpus_path(corpus_name)
        if corpus_path is None:
            return ToolResult(
                tool_call_id=tool_call.id,
                content="",
                error=f"Unknown corpus: {corpus_name}. Available: {AVAILABLE_CORPORA}",
            )
    
    if not corpus_path.exists():
        return ToolResult(
            tool_call_id=tool_call.id,
            content="",
            error=f"Corpus '{corpus_name}' not downloaded. Run: wafer corpus download {corpus_name}",
        )
    
    # Search
    results = _search_files(corpus_path, query, max_results)
    
    if not results:
        return ToolResult(
            tool_call_id=tool_call.id,
            content=f"No results found for query: {query}",
            error=None,
        )
    
    # Format output
    output_parts = [f"Found {len(results)} results for: {query}\n"]
    
    for i, result in enumerate(results, 1):
        output_parts.append(f"\n{'='*60}")
        output_parts.append(f"[{i}] {result['file']} (score: {result['score']})")
        output_parts.append("=" * 60)
        
        for snippet in result["snippets"]:
            output_parts.append(snippet)
            output_parts.append("-" * 40)
    
    return ToolResult(
        tool_call_id=tool_call.id,
        content="\n".join(output_parts),
        error=None,
    )
