#!/usr/bin/env python3
"""
MCP server for full-text search using Tantivy.
Initializes with a fixed JSONL file and provides search/read tools via HTTP.
"""

import argparse
import asyncio
import json
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Literal, Optional, TypeAlias, TypedDict, Union

from fastmcp import FastMCP
from pydantic import BaseModel

from ..types import SearchResultContent
from ..filters import filter_document_content
from ..index import TantivySearch
from ..core import (
    IndexConfig,
    load_config_from_yaml,
    load_jsonl_data,
    ensure_id_column,
    preview_text,
    normalize_field_value,
    normalize_metadata_value,
)


class SearchIndexDict(TypedDict):
    config: IndexConfig
    search_index: TantivySearch
    original_data: list[dict]
    index_path: str


class FullTextSearchMCP:
    """Encapsulates full-text search MCP server functionality with support for multiple indices."""

    def __init__(
        self,
        server_name: str,
        server_description: str,
        indices: Optional[list[IndexConfig]] = None,
        max_workers: int = 4,
    ):
        """Initialize the MCP server.

        Args:
            server_name: Name of the server
            server_description: Description of the server
            indices: List of index configurations
            max_workers: Number of worker threads for search operations (default: 4)
        """
        self.server_name = server_name
        self.server_description = server_description
        # index_name -> {config, search_index, original_data}
        self.indices: dict[str, SearchIndexDict] = {}
        self.citations = {}
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        if indices:
            for index_config in indices:
                self.add_index(index_config)

    def __del__(self):
        """Cleanup executor on object destruction."""
        if hasattr(self, "executor"):
            self.executor.shutdown(wait=True)

    def add_index(self, config: IndexConfig) -> None:
        """Add and initialize a single index."""
        print(f"Adding index '{config.name}'...")

        # Load data
        print(f"Loading data from {config.data_file}...")
        original_data = load_jsonl_data(config.data_file)
        if not original_data:
            raise ValueError(f"No data found in file {config.data_file}")

        # Ensure id column exists (auto-generate if missing and id_column is 'id')
        original_data = ensure_id_column(original_data, config.id_column)

        metadata_columns = config.metadata_columns or []
        metadata_set = set(metadata_columns)

        # Enforce disjoint searchable vs metadata columns
        reserved = set(config.searchable_columns) | {
            config.text_column,
            config.id_column,
        }
        overlap = metadata_set & reserved
        if overlap:
            raise ValueError(
                "metadata_columns must be disjoint from searchable/text/id columns. "
                f"Overlap: {sorted(overlap)}"
            )

        # Validate columns exist
        sample_record = original_data[0]
        missing_cols = []
        if config.id_column not in sample_record:
            missing_cols.append(config.id_column)
        if config.text_column not in sample_record:
            missing_cols.append(config.text_column)
        for col in config.searchable_columns:
            if col not in sample_record:
                missing_cols.append(col)
        for col in metadata_columns:
            if col not in sample_record:
                missing_cols.append(col)

        if missing_cols:
            raise ValueError(
                f"Missing columns in data for index '{config.name}': {missing_cols}"
            )

        # Set default index path if not provided
        index_path = config.index_path
        if index_path is None:
            data_file_path = Path(config.data_file)
            index_path = str(data_file_path.parent / f"{data_file_path.stem}_index")

        # Create search index
        print(f"Building search index at {index_path}...")
        search_index = TantivySearch(
            index_path,
            include_fields=config.include_fields,
            exclude_fields=config.exclude_fields,
            metadata_fields=metadata_columns,
        )

        # Build index with all searchable columns, metadata columns, plus id
        all_columns = (
            [config.id_column, config.text_column]
            + [
                col
                for col in config.searchable_columns
                if col not in [config.id_column, config.text_column]
            ]
            + metadata_columns
        )

        # Filter records to only include relevant columns and normalize values
        filtered_records = []
        for record in original_data:
            filtered_record = {}
            for col in all_columns:
                if col in record:
                    if col in metadata_set:
                        normalized_values = normalize_metadata_value(record[col])
                        if not normalized_values:
                            continue
                        if len(normalized_values) == 1:
                            filtered_record[col] = normalized_values[0]
                        else:
                            filtered_record[col] = normalized_values
                    else:
                        # Normalize field values (handles JSON/list columns)
                        filtered_record[col] = normalize_field_value(record[col])

            # Ensure id is string
            filtered_record["id"] = str(filtered_record.get(config.id_column, ""))
            # Remove the original id column if it's different from "id"
            if config.id_column != "id" and config.id_column in filtered_record:
                del filtered_record[config.id_column]
            filtered_records.append(filtered_record)

        # Build index without deduplication
        search_index.build_index(filtered_records, deduplicate_strategy=None)

        # Cache original documents for efficient fetching
        search_index.cache_documents_for_fetch(original_data, config.id_column)

        # Store the index data
        self.indices[config.name] = {
            "config": config,
            "search_index": search_index,
            "original_data": original_data,
            "index_path": index_path,
        }

        # Add all documents to citations dict as index_name:doc_id -> {title, url}
        for record in original_data:
            doc_id = str(record.get(config.id_column, ""))
            citation_key = f"{config.name}:{doc_id}"
            citation_value = {}

            # Add title if available
            if "title" in record:
                citation_value["title"] = record["title"]

            # Add url if available
            if "url" in record:
                citation_value["url"] = record["url"]

            # Only add to citations if we have title or url
            if citation_value:
                self.citations[citation_key] = citation_value

        print(
            f"Successfully initialized index '{config.name}' with {len(original_data)} documents"
        )

    async def search(
        self,
        queries: list[str],
        index_name: str = "default",
        limit: int = 5,
        metadata_filters: dict[str, list[str]] | None = None,
    ) -> str:
        """Search documents using multiple queries with result fusion and return previews."""
        if index_name not in self.indices:
            available = list(self.indices.keys())
            return (
                f"Error: Index '{index_name}' not found. Available indices: {available}"
            )

        index_data = self.indices[index_name]
        search_index = index_data["search_index"]
        config = index_data["config"]

        if metadata_filters:
            if not config.metadata_columns:
                return "Error: Metadata filters requested but no metadata fields configured"
            unknown_fields = set(metadata_filters) - set(config.metadata_columns)
            if unknown_fields:
                return "Error: Unknown metadata filter fields: " + ", ".join(
                    sorted(unknown_fields)
                )

        try:
            # Run blocking search in thread pool
            loop = asyncio.get_event_loop()

            def _search():
                # Search across all searchable columns
                all_searchable = [config.text_column] + config.searchable_columns
                # Remove duplicates while preserving order
                search_fields = []
                for field in all_searchable:
                    if field not in search_fields and field != config.id_column:
                        search_fields.append(field)

                return search_index.search(
                    queries,
                    search_fields,
                    limit=limit,
                    metadata_filters=metadata_filters,
                )

            results = await loop.run_in_executor(self.executor, _search)

            # Add line count metadata to results
            for result in results:
                text_content = result.content.get(config.text_column, "")
                if text_content:
                    result.content["_line_count"] = text_content.count("\n") + 1  # type: ignore
                    result.content["_content_length"] = len(text_content)  # type: ignore
                else:
                    result.content["_line_count"] = 0  # type: ignore
                    result.content["_content_length"] = 0  # type: ignore

            if not results:
                query_str = "', '".join(queries)
                return f"No results found for queries: ['{query_str}'] in index '{index_name}'"

            # Format results with previews
            query_str = "', '".join(queries)
            formatted_results = [
                f"Found {len(results)} results for queries ['{query_str}'] in index '{index_name}':\n\n"
            ]

            for i, result in enumerate(results, 1):
                content: SearchResultContent = result.content  # type: ignore
                doc_id = result.id
                score = result.score

                # Get preview of main text content
                main_text = content.get(config.text_column, "")
                preview = preview_text(main_text)

                # add keywords
                formatted_results.append(f"Result {i} (Score: {score:.3f})")
                formatted_results.append(f"ID: {doc_id}")
                if "title" in content:
                    title = content["title"]
                    formatted_results.append(f"Title: {title}")
                if "description" in content:
                    description = content["description"]
                    formatted_results.append(
                        f"Description: {preview_text(description)}"
                    )
                if "keywords" in content:
                    keywords = content["keywords"]
                    keywords_lst = keywords.split(", ")
                    if keywords_lst:
                        top10_keywords = keywords_lst[:12]
                        remaining_keywords = max(len(keywords_lst) - 10, 0)
                        keywords_formatted = ", ".join(top10_keywords)
                        if remaining_keywords > 0:
                            keywords_formatted += f" (+ {remaining_keywords} more)"
                        formatted_results.append(f"Keywords: {keywords_formatted}")
                formatted_results.append(f"Lines: {content.get('_line_count', 0)}")
                formatted_results.append(f"Content (preview): {preview}")

                formatted_results.append("")  # Empty line between results

            return "\n".join(formatted_results)

        except Exception as e:
            return f"Error during search in index '{index_name}': {str(e)}"

    async def read_documents(
        self,
        document_ids: list[str],
        index_name: str = "default",
        filters: list[dict] | None = None,
    ) -> str:
        """Retrieve full content for specific document IDs with optional filtering.

        Args:
            document_ids: List of document IDs to retrieve
            index_name: Name of the index to search in
            filters: Optional list of filter specs (one per document). Each filter can be:
                - Range: Filter by line range (start_line, end_line)
                - GrepFilter: Filter by regex pattern (grep_pattern, context_lines)
                - dict: Plain dict with filter spec (for direct Python calls)
                - None: No filtering for this document
        """
        if index_name not in self.indices:
            available = list(self.indices.keys())
            return (
                f"Error: Index '{index_name}' not found. Available indices: {available}"
            )

        # Validate filters length matches document_ids length
        if filters is not None and len(filters) != len(document_ids):
            return f"Error: filters length ({len(filters)}) must match document_ids length ({len(document_ids)})"

        index_data = self.indices[index_name]
        search_index = index_data["search_index"]
        config = index_data["config"]

        try:
            # Run in thread pool for consistency
            loop = asyncio.get_event_loop()

            def _fetch():
                # Use TantivySearch's cached fetch method
                return search_index.fetch(document_ids)

            results, found_ids, missing_ids = await loop.run_in_executor(
                self.executor, _fetch
            )

            # Format results
            if not results:
                return f"No documents found for IDs: {document_ids} in index '{index_name}'"

            formatted_results = [
                f"Retrieved {len(results)} documents from index '{index_name}':\n"
            ]

            for i, record in enumerate(results):
                doc_id = found_ids[i]
                formatted_results.append(f"Document {i + 1} (ID: {doc_id})")
                formatted_results.append("-" * 3)

                # Apply filtering if specified
                filter_obj = filters[i] if filters else None
                # filter_to_dict(filter_obj) if filter_obj else None
                filter_spec = filter_obj

                # Show all fields
                for key, value in record.items():
                    # Apply filtering to the main text column
                    if filter_spec and key == config.text_column:
                        filtered_value = filter_document_content(
                            str(value), filter_spec
                        )
                        formatted_results.append(f"{key}: {filtered_value}")
                    else:
                        formatted_results.append(f"{key}: {value}")

                formatted_results.append("")  # Empty line between documents

            if missing_ids:
                formatted_results.append(f"Missing document IDs: {missing_ids}")

            return "\n".join(formatted_results)

        except Exception as e:
            return f"Error retrieving documents from index '{index_name}': {str(e)}"

    def create_mcp_server(self) -> FastMCP:
        """Create MCP server with tools."""
        mcp = FastMCP(self.server_name)

        # Create search and read tools
        search_tool_name = "search"
        search_description = (
            "Search documents in one of the available search indices using one or more queries, returning return previews. "
            "Search uses full-text keyword search with Tantivy (not neural/embedding search), "
            "so longer queries or multiple queries with synonyms and more keywords can improve recall. "
            "Multi-queries will have their results fused with reciprocal rank fusion. "
            "Use multiple queries + higher limit for speed/breadth of search, since some queries may not return any results "
            "if you don't know what the index contains.\n\n"
            "Available indices:\n"
        )
        for index_name, index_data in self.indices.items():
            config = index_data["config"]
            search_description += f" - {index_name}: {config.description}\n"

        search_description += (
            "Args:\n\t\n\tindex (str): which index to search.\n\tqueries: List of search query strings (allows multiple queries for better results)"
            "\n\tlimit: Maximum total number of results to return (default 5)"
            "\n\tmetadata_filters: Optional metadata filters {field: [value1, value2]} (OR within field, AND across fields)"
        )

        fetch_tool_name = "fetch"
        fetch_description = (
            "Retrieve full content from an index for specific document IDs from that index. "
            "Each document can be optionally filtered to read specific portions.\n\n"
            "Args:\n\tindex (str): which index to read documents from."
            "\n\tdocument_ids: List of document IDs to retrieve"
            "\n\tfilters (optional): List of filter specs (one per document). "
            "Note that the schema technically allows this to be a string, but that's "
            "for compatibility, if not null you should return a structured list[dict] as follows. "
            "You can either:\n"
            "\t\t- Filter lines by including start_line (int) and end_line (int), OR\n"
            "\t\t- Filter by a regex by including grep_pattern (str) and context_lines (int), "
            "which will show all lines with the pattern + context_lines lines of context, OR\n"
            "\t\t- Read the entire document by including an empty filter object {}.\n\n"
            "Examples:\n"
            '\t[{"start_line": 0, "end_line": 50}] # First 50 lines of first doc\n'
            '\t[{"grep_pattern": "section", "context_lines": 2}]  # Lines matching "section" with 2 lines context\n'
            '\t[{}, {"start_line": 100, "end_line": 200}]  # Full first doc, lines 100-200 of second doc\n'
        )

        batch_tool_name = "batch"
        batch_description = (
            "Execute any number of search/fetch operations in parallel. "
            "This is recommended to speed up your research, reducing "
            "latency for the end user. In order to use this tool, "
            "simply provide a list (array) of commands, where each command "
            'has name ("search" or "fetch"), index, and either queries + limit (for search) '
            "or document_ids (for fetch). For example:\n\n"
            '[\n\t{"name": "search", "index": "my_index", "queries": ["query1", "query2"], "limit": 5},\n'
            '\t{"name": "fetch", "index": "my_index", "document_ids": ["121", "3"]}, "filters": [{}, {}]\n]\n\n'
            "For example, you can read one set of documents while performing your "
            "next search in parallel; or perform several searches in parallel across several indices."
        )

        async def search_tool(
            index: str,
            queries: list[str],
            limit: int = 5,
            metadata_filters: dict[str, list[str]] | None = None,
        ) -> str:
            return await self.search(queries, index, limit, metadata_filters)

        async def fetch_tool(
            index: str,
            document_ids: list[str],
            filters: list[dict] | str | None = None,
        ) -> str:
            if filters is None:
                return await self.read_documents(document_ids, index, filters)
            else:
                parsed_filters = []
                if isinstance(filters, str):
                    filters = json.loads(filters)
                assert isinstance(filters, list), "filters should be a list"
                for filter in filters:
                    # determine the type
                    if isinstance(filter, str):
                        filter = json.loads(filter)
                    if len(filter) == 0:
                        parsed_filters.append({})
                    elif "start_line" in filter and "end_line" in filter:
                        parsed_filters.append(filter)
                    elif "grep_pattern" in filter and "context_lines" in filter:
                        parsed_filters.append(filter)
                    else:
                        print(
                            f"WARNING: Got invalid filter - {filter}. Will return full doc."
                        )
                        parsed_filters.append({})

                return await self.read_documents(document_ids, index, parsed_filters)

        class SearchCall(BaseModel):
            name: Literal["search"]
            index: str
            queries: list[str]
            limit: int = 5
            metadata_filters: dict[str, list[str]] | None = None

        class FetchCall(BaseModel):
            name: Literal["fetch"]
            index: str
            document_ids: list[str]
            filters: list[dict] | str | None = None

        Command: TypeAlias = Union[SearchCall, FetchCall]  # type: ignore

        async def batch_tool(commands: list[Command]) -> str:
            # Execute all commands concurrently
            tasks = []
            for command in commands:
                if command.name == "search":
                    task = search_tool(
                        command.index,
                        command.queries,
                        command.limit,
                        command.metadata_filters,
                    )
                elif command.name == "fetch":
                    task = fetch_tool(
                        command.index, command.document_ids, command.filters
                    )
                else:
                    # Create a coroutine that returns an error message
                    async def error_task():
                        return "unknown command: " + command.name

                    task = error_task()
                tasks.append(task)

            # Wait for all tasks to complete
            results = await asyncio.gather(*tasks, return_exceptions=True)

            # Convert exceptions to error messages
            formatted_results = []
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    formatted_results.append(
                        f"error executing {commands[i].name}: {str(result)}"
                    )
                else:
                    formatted_results.append(result)

            return "# Batch Tool Results\n\n" + "\n\n---\n\n".join(formatted_results)

        # Register the tools
        mcp.tool(name=search_tool_name, description=search_description)(search_tool)
        mcp.tool(name=fetch_tool_name, description=fetch_description)(fetch_tool)
        mcp.tool(name=batch_tool_name, description=batch_description)(batch_tool)

        # add resource for citations
        def get_citation(index: str, doc_id: str):
            key = f"{index}:{doc_id}"
            if key in self.citations:
                return self.citations[key]
            else:
                return None

        mcp.resource("cite://{index}/{doc_id}")(get_citation)
        return mcp

    def run_server(self, host: str = "0.0.0.0", port: int = 8000) -> None:
        """Run the MCP server."""
        mcp = self.create_mcp_server()
        print(f"Starting MCP server on {host}:{port}")
        print(f"MCP endpoint: http://{host}:{port}/mcp")
        mcp.run(
            transport="streamable-http",
            host=host,
            port=port,
            stateless_http=True,
        )


def main():
    """CLI entry point for running the MCP server."""
    parser = argparse.ArgumentParser(description="Full-text search MCP server")

    # Add config file option
    parser.add_argument(
        "--config",
        help="Path to YAML configuration file. If provided, other arguments are ignored.",
    )

    # Required arguments for server configuration (when not using config file)
    parser.add_argument(
        "--name",
        help="Name for the server (should be alpha + underscores)",
    )
    parser.add_argument(
        "--description",
        help="Description of what this search index contains",
    )
    parser.add_argument("--data-file", help="Path to JSONL file containing documents")
    parser.add_argument(
        "--id-column",
        help="Name of column containing unique document IDs",
    )
    parser.add_argument("--text-column", help="Name of main text column for content")
    parser.add_argument(
        "--searchable-columns",
        nargs="+",
        help="List of column names to make searchable",
    )

    # Optional arguments
    parser.add_argument(
        "--index-path", help="Path for index storage (defaults to data_file_index)"
    )
    parser.add_argument("--host", default="0.0.0.0", help="Host for HTTP server")
    parser.add_argument("--port", type=int, default=8000, help="Port for HTTP server")
    parser.add_argument(
        "--max-workers",
        type=int,
        default=4,
        help="Number of worker threads for concurrent search operations (default: 4)",
    )

    args = parser.parse_args()

    try:
        if args.config:
            # Load configuration from YAML file
            server_name, server_description, index_configs, server_config = (  # type: ignore
                load_config_from_yaml(args.config)
            )

            # Create multi-index search server with configurable max_workers
            max_workers = server_config.get("max_workers", args.max_workers)
            search_server = FullTextSearchMCP(
                server_name=server_name,
                server_description=server_description,
                indices=index_configs,
                max_workers=max_workers,
            )

            # Use host/port from YAML config if specified, otherwise use CLI args
            host = server_config.get("host", args.host)
            port = server_config.get("port", args.port)

        else:
            raise NotImplementedError("Must use YAML config to launch server.")

        # Start HTTP server
        search_server.run_server(host=host, port=port)

    except Exception as e:
        print(f"Error starting server: {e}", file=sys.stderr)
        return 1

    return 0


if __name__ == "__main__":
    main()
