"""
Core search service implementation.

This module provides the SearchService class which encapsulates all search logic
independent of HTTP framework. It can be used by Lambda, FastAPI, or any other
deployment pattern.
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Any, TypedDict

from .core import (
    IndexConfig,
    BatchCommand,
    SearchResponse,
    FetchResponse,
    BatchResponse,
    IndexInfo,
    IndicesResponse,
    HealthResponse,
    load_jsonl_data,
    ensure_id_column,
    preview_text,
    normalize_field_value,
    normalize_metadata_value,
)
from .filters import filter_document_content
from .index import TantivySearch


class IndexData(TypedDict):
    """Type definition for stored index data."""

    config: IndexConfig
    search_index: TantivySearch
    original_data: list[dict]
    index_path: str


@dataclass
class SearchServiceConfig:
    """Configuration for SearchService initialization."""

    server_name: str = "search_service"
    server_description: str = "Full-text search service"


class SearchService:
    """Core search service - used by both Lambda and HTTP API servers.

    This class encapsulates all search logic independent of the HTTP framework,
    enabling code reuse across different deployment patterns (Lambda, FastAPI, etc.).
    """

    def __init__(self, config: SearchServiceConfig | None = None):
        """Initialize the search service.

        Args:
            config: Optional service configuration
        """
        self.config = config or SearchServiceConfig()
        self.indices: dict[str, IndexData] = {}
        self.citations: dict[str, dict[str, str]] = {}

    def add_index(
        self,
        index_config: IndexConfig,
        data: list[dict] | None = None,
    ) -> None:
        """Add and initialize a single index.

        Args:
            index_config: Configuration for the index
            data: Pre-loaded data (if None, loads from index_config.data_file)
        """
        print(f"Adding index '{index_config.name}'...")

        # Load data if not provided
        if data is None:
            print(f"Loading data from {index_config.data_file}...")
            data = load_jsonl_data(index_config.data_file)

        if not data:
            raise ValueError(f"No data found for index '{index_config.name}'")

        # Ensure id column exists
        data = ensure_id_column(data, index_config.id_column)

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

        # Enforce disjoint searchable vs metadata columns
        reserved = set(index_config.searchable_columns) | {
            index_config.text_column,
            index_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 = data[0]
        missing_cols = []
        if index_config.id_column not in sample_record:
            missing_cols.append(index_config.id_column)
        if index_config.text_column not in sample_record:
            missing_cols.append(index_config.text_column)
        for col in index_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 '{index_config.name}': {missing_cols}"
            )

        # Set default index path if not provided
        index_path = index_config.index_path
        if index_path is None:
            data_file_path = Path(index_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=index_config.include_fields,
            exclude_fields=index_config.exclude_fields,
            metadata_fields=metadata_columns,
        )

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

        # Filter records to only include relevant columns and normalize values
        filtered_records = []
        for record in 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:
                        filtered_record[col] = normalize_field_value(record[col])

            # Ensure id is string and use 'id' as key
            filtered_record["id"] = str(filtered_record.get(index_config.id_column, ""))
            if (
                index_config.id_column != "id"
                and index_config.id_column in filtered_record
            ):
                del filtered_record[index_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(data, index_config.id_column)

        # Store the index data
        self.indices[index_config.name] = IndexData(
            config=index_config,
            search_index=search_index,
            original_data=data,
            index_path=index_path,
        )

        # Add all documents to citations dict
        for record in data:
            doc_id = str(record.get(index_config.id_column, ""))
            citation_key = f"{index_config.name}:{doc_id}"
            citation_value = {}

            if "title" in record:
                citation_value["title"] = record["title"]
            if "url" in record:
                citation_value["url"] = record["url"]

            if citation_value:
                self.citations[citation_key] = citation_value

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

    def search(
        self,
        index: str,
        queries: list[str],
        limit: int = 5,
        include_content: bool = False,
        metadata_filters: dict[str, list[str]] | None = None,
    ) -> SearchResponse:
        """Search documents in an index.

        Args:
            index: Name of the index to search
            queries: List of search query strings (fused with reciprocal rank fusion)
            limit: Maximum number of results to return
            include_content: If True, include full content dict in results
            metadata_filters: Optional metadata filters to apply

        Returns:
            SearchResponse with results

        Raises:
            ValueError: If index not found
        """
        if index not in self.indices:
            raise ValueError(
                f"Index '{index}' not found. Available: {list(self.indices.keys())}"
            )

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

        if metadata_filters:
            if not config.metadata_columns:
                raise ValueError(
                    f"Index '{index}' does not support metadata filters (no metadata_columns configured)"
                )
            unknown_fields = set(metadata_filters) - set(config.metadata_columns)
            if unknown_fields:
                raise ValueError(
                    f"Unknown metadata filter fields for index '{index}': {sorted(unknown_fields)}"
                )

        # Build search fields
        all_searchable = [config.text_column] + config.searchable_columns
        search_fields = []
        for field in all_searchable:
            if field not in search_fields and field != config.id_column:
                search_fields.append(field)

        # Execute search
        results = search_index.search(
            queries, search_fields, limit=limit, metadata_filters=metadata_filters
        )

        # Format results
        formatted_results = []
        for result in results:
            content = result.content
            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)

            # Build result dict
            result_dict: dict[str, Any] = {
                "id": doc_id,
                "score": score,
                "title": content.get("title"),
                "description": preview_text(content.get("description", ""), 50)
                if content.get("description")
                else None,
                "preview": preview,
            }

            if include_content:
                result_dict["content"] = content

            formatted_results.append(result_dict)

        return SearchResponse(
            index=index,
            queries=queries,
            results=formatted_results,
            total=len(formatted_results),
        )

    def fetch(
        self,
        index: str,
        document_ids: list[str],
        filters: list[dict[str, Any]] | None = None,
    ) -> FetchResponse:
        """Fetch full documents by ID.

        Args:
            index: Name of the index to fetch from
            document_ids: List of document IDs to retrieve
            filters: Optional list of filter specs (one per document)

        Returns:
            FetchResponse with documents

        Raises:
            ValueError: If index not found or filters length mismatch
        """
        if index not in self.indices:
            raise ValueError(
                f"Index '{index}' not found. Available: {list(self.indices.keys())}"
            )

        if filters is not None and len(filters) != len(document_ids):
            raise ValueError(
                f"filters length ({len(filters)}) must match document_ids length ({len(document_ids)})"
            )

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

        # Fetch documents
        results, found_ids, missing_ids = search_index.fetch(document_ids)

        # Apply filters if provided
        if filters:
            for i, record in enumerate(results):
                filter_spec = filters[i] if i < len(filters) else None
                if filter_spec and config.text_column in record:
                    filtered_content = filter_document_content(
                        str(record[config.text_column]), filter_spec
                    )
                    record[config.text_column] = filtered_content

        return FetchResponse(
            index=index,
            documents=results,
            found_ids=found_ids,
            missing_ids=missing_ids,
        )

    def batch(self, commands: list[BatchCommand]) -> BatchResponse:
        """Execute multiple search/fetch operations.

        Note: This is a synchronous implementation. For async execution,
        use the async wrapper in the HTTP server.

        Args:
            commands: List of commands to execute

        Returns:
            BatchResponse with results for each command
        """
        results = []
        for cmd in commands:
            try:
                if cmd.name == "search":
                    if not cmd.queries:
                        results.append({"error": "queries required for search"})
                        continue
                    result = self.search(
                        cmd.index,
                        cmd.queries,
                        cmd.limit,
                        cmd.include_content,
                        cmd.metadata_filters,
                    )
                    results.append(result.model_dump())
                elif cmd.name == "fetch":
                    if not cmd.document_ids:
                        results.append({"error": "document_ids required for fetch"})
                        continue
                    result = self.fetch(cmd.index, cmd.document_ids, cmd.filters)
                    results.append(result.model_dump())
                else:
                    results.append({"error": f"Unknown command: {cmd.name}"})
            except ValueError as e:
                results.append({"error": str(e)})
            except Exception as e:
                results.append({"error": str(e)})

        return BatchResponse(results=results)

    def list_indices(self) -> IndicesResponse:
        """List all available indices.

        Returns:
            IndicesResponse with index information
        """
        indices_info = []
        for name, data in self.indices.items():
            config = data["config"]
            indices_info.append(
                IndexInfo(
                    name=name,
                    description=config.description,
                    document_count=len(data["original_data"]),
                    searchable_columns=config.searchable_columns,
                    metadata_columns=config.metadata_columns,
                )
            )
        return IndicesResponse(indices=indices_info)

    def health_check(self) -> HealthResponse:
        """Check service health.

        Returns:
            HealthResponse with status and loaded indices
        """
        return HealthResponse(
            status="healthy",
            indices=list(self.indices.keys()),
            server=self.config.server_name,
        )

    def get_citation(self, index: str, doc_id: str) -> dict[str, str] | None:
        """Get citation metadata for a document.

        Args:
            index: Index name
            doc_id: Document ID

        Returns:
            Citation dict with title/url if found, None otherwise
        """
        key = f"{index}:{doc_id}"
        return self.citations.get(key)


__all__ = [
    "SearchService",
    "SearchServiceConfig",
    "IndexData",
]
