# fitz_ai/retrieval/sparse/index.py
"""
Sparse Index for hybrid retrieval using PostgreSQL tsvector.

Uses PostgreSQL's built-in full-text search capabilities for keyword-based retrieval.
Complements dense (semantic) search with exact keyword matching.

The tsvector index is automatically maintained by the pgvector chunks table,
so this module provides a query interface that leverages that index.
"""

from __future__ import annotations

from dataclasses import dataclass

from fitz_ai.logging.logger import get_logger
from fitz_ai.storage import get_connection_manager

logger = get_logger(__name__)


@dataclass
class SparseHit:
    """A single sparse search result."""

    chunk_id: str
    score: float


class SparseIndex:
    """
    PostgreSQL tsvector-based sparse index for hybrid retrieval.

    Uses the existing tsvector index on the pgvector chunks table.
    No separate index building required - the tsvector is automatically
    generated and indexed when chunks are inserted.

    Usage (query):
        index = SparseIndex(collection="my_collection")
        hits = index.search("hello world", k=10)
    """

    def __init__(self, collection: str):
        """
        Initialize sparse index for a collection.

        Args:
            collection: Collection name (used for database routing)
        """
        self.collection = collection
        self._manager = get_connection_manager()
        self._manager.start()

    @classmethod
    def load(cls, collection: str) -> "SparseIndex":
        """
        Load a sparse index for a collection.

        For PostgreSQL tsvector, this just creates an instance
        since the index is maintained automatically by the database.

        Args:
            collection: Collection name

        Returns:
            SparseIndex instance
        """
        return cls(collection=collection)

    def build(self, chunk_ids: list[str], contents: list[str]) -> None:
        """
        Build the sparse index from chunks.

        For PostgreSQL tsvector, this is a no-op since the tsvector
        column is automatically generated when chunks are inserted
        via the pgvector upsert.

        Args:
            chunk_ids: List of chunk IDs (ignored)
            contents: List of chunk contents (ignored)
        """
        # No-op: tsvector is auto-generated by PostgreSQL
        logger.debug(
            f"Sparse index build called for '{self.collection}' - "
            "tsvector is auto-maintained by PostgreSQL"
        )

    def add(self, chunk_ids: list[str], contents: list[str]) -> None:
        """
        Add new chunks to the index.

        For PostgreSQL tsvector, this is a no-op since the tsvector
        column is automatically updated when chunks are upserted.

        Args:
            chunk_ids: New chunk IDs (ignored)
            contents: New chunk contents (ignored)
        """
        # No-op: tsvector is auto-maintained by PostgreSQL
        pass

    def save(self) -> None:
        """
        Save the sparse index.

        For PostgreSQL tsvector, this is a no-op since the index
        is persisted automatically in the database.
        """
        # No-op: PostgreSQL handles persistence
        pass

    def search(self, query: str, k: int = 10) -> list[SparseHit]:
        """
        Search the sparse index using PostgreSQL full-text search.

        Args:
            query: Query string
            k: Number of results to return

        Returns:
            List of SparseHit with chunk_id and score
        """
        if not query or not query.strip():
            return []

        try:
            with self._manager.connection(self.collection) as conn:
                # Check if chunks table exists
                table_exists = conn.execute(
                    """
                    SELECT EXISTS (
                        SELECT FROM information_schema.tables
                        WHERE table_name = 'chunks'
                    )
                    """
                ).fetchone()[0]

                if not table_exists:
                    return []

                # Use plainto_tsquery for natural language queries
                # ts_rank_cd provides a relevance score
                cursor = conn.execute(
                    """
                    SELECT id, ts_rank_cd(content_tsv, plainto_tsquery('english', %s)) as score
                    FROM chunks
                    WHERE content_tsv @@ plainto_tsquery('english', %s)
                    ORDER BY score DESC
                    LIMIT %s
                    """,
                    (query, query, k),
                )

                return [
                    SparseHit(chunk_id=row[0], score=float(row[1]) if row[1] else 0.0)
                    for row in cursor.fetchall()
                ]

        except Exception as e:
            logger.warning(f"Sparse search failed: {e}")
            return []

    def search_phrase(self, query: str, k: int = 10) -> list[SparseHit]:
        """
        Search using phrase matching (words must appear together).

        Args:
            query: Query string (treated as a phrase)
            k: Number of results to return

        Returns:
            List of SparseHit with chunk_id and score
        """
        if not query or not query.strip():
            return []

        try:
            with self._manager.connection(self.collection) as conn:
                # Check if chunks table exists
                table_exists = conn.execute(
                    """
                    SELECT EXISTS (
                        SELECT FROM information_schema.tables
                        WHERE table_name = 'chunks'
                    )
                    """
                ).fetchone()[0]

                if not table_exists:
                    return []

                # Use phraseto_tsquery for phrase matching
                cursor = conn.execute(
                    """
                    SELECT id, ts_rank_cd(content_tsv, phraseto_tsquery('english', %s)) as score
                    FROM chunks
                    WHERE content_tsv @@ phraseto_tsquery('english', %s)
                    ORDER BY score DESC
                    LIMIT %s
                    """,
                    (query, query, k),
                )

                return [
                    SparseHit(chunk_id=row[0], score=float(row[1]) if row[1] else 0.0)
                    for row in cursor.fetchall()
                ]

        except Exception as e:
            logger.warning(f"Phrase search failed: {e}")
            return []

    def search_websearch(self, query: str, k: int = 10) -> list[SparseHit]:
        """
        Search using websearch syntax (supports AND, OR, NOT, quotes).

        Example queries:
        - "machine learning" (phrase)
        - python AND django (both required)
        - python OR ruby (either)
        - python -java (exclude java)

        Args:
            query: Query string in websearch syntax
            k: Number of results to return

        Returns:
            List of SparseHit with chunk_id and score
        """
        if not query or not query.strip():
            return []

        try:
            with self._manager.connection(self.collection) as conn:
                # Check if chunks table exists
                table_exists = conn.execute(
                    """
                    SELECT EXISTS (
                        SELECT FROM information_schema.tables
                        WHERE table_name = 'chunks'
                    )
                    """
                ).fetchone()[0]

                if not table_exists:
                    return []

                # Use websearch_to_tsquery for advanced syntax
                cursor = conn.execute(
                    """
                    SELECT id, ts_rank_cd(content_tsv, websearch_to_tsquery('english', %s)) as score
                    FROM chunks
                    WHERE content_tsv @@ websearch_to_tsquery('english', %s)
                    ORDER BY score DESC
                    LIMIT %s
                    """,
                    (query, query, k),
                )

                return [
                    SparseHit(chunk_id=row[0], score=float(row[1]) if row[1] else 0.0)
                    for row in cursor.fetchall()
                ]

        except Exception as e:
            logger.warning(f"Websearch failed: {e}")
            return []

    def is_ready(self) -> bool:
        """Check if index is ready for querying."""
        try:
            with self._manager.connection(self.collection) as conn:
                result = conn.execute(
                    """
                    SELECT EXISTS (
                        SELECT FROM information_schema.tables
                        WHERE table_name = 'chunks'
                    )
                    """
                ).fetchone()
                return result[0] if result else False
        except Exception:
            return False

    def __len__(self) -> int:
        """Number of indexed documents."""
        try:
            with self._manager.connection(self.collection) as conn:
                result = conn.execute("SELECT COUNT(*) FROM chunks").fetchone()
                return result[0] if result else 0
        except Exception:
            return 0

    # Backwards compatibility properties
    vectorizer = None
    doc_vectors = None
    chunk_ids: list[str] = []
