"""Vector search and FTS operations mixin."""
from __future__ import annotations

import re
from typing import Any, List, Optional

import aiosqlite

from cortex.storage._helpers import _cosine_similarity, _pack_floats, _unpack_floats
from cortex.types import FTSResult, VectorResult


class ClaimVectorMixin:
    """Mixin providing vector search (sqlite-vec + brute-force) and FTS."""

    async def _vec_search_sqlite_vec(
        self,
        conn: aiosqlite.Connection,
        query_embedding: List[float],
        item_type: Optional[str],
        top_k: int,
        owner_id: str,
    ) -> List[VectorResult]:
        """sqlite-vec powered KNN search via vec_index virtual table."""
        dim = len(query_embedding)
        blob = _pack_floats(query_embedding)

        try:
            if item_type:
                async with conn.execute(
                    """
                    SELECT v.item_type, v.item_id, v.distance
                    FROM vec_index v
                    WHERE v.item_type = ?
                      AND v.embedding MATCH ?
                      AND k = ?
                    ORDER BY v.distance
                    """,
                    (item_type, blob, top_k),
                ) as cur:
                    rows = await cur.fetchall()
            else:
                async with conn.execute(
                    """
                    SELECT v.item_type, v.item_id, v.distance
                    FROM vec_index v
                    WHERE v.embedding MATCH ?
                      AND k = ?
                    ORDER BY v.distance
                    """,
                    (blob, top_k),
                ) as cur:
                    rows = await cur.fetchall()

            # Enforce owner isolation for claim embeddings. sqlite-vec does not
            # support arbitrary joins in the KNN query itself, so we post-filter.
            if item_type == "claim" or item_type is None:
                claim_ids = [r["item_id"] for r in rows if r["item_type"] == "claim"]
                if claim_ids:
                    placeholders = ",".join("?" for _ in claim_ids)
                    async with conn.execute(
                        f"SELECT id FROM semantic_claim WHERE owner_id = ? AND id IN ({placeholders})"
                        " AND status = 'active'"
                        " AND (is_deleted = 0 OR is_deleted IS NULL)",
                        [owner_id, *claim_ids],
                    ) as cur:
                        allowed = {row["id"] for row in await cur.fetchall()}
                    rows = [
                        r for r in rows
                        if r["item_type"] != "claim" or r["item_id"] in allowed
                    ]

            results = [
                VectorResult(
                    item_type=r["item_type"],
                    item_id=r["item_id"],
                    score=1.0 - r["distance"],  # convert L2 distance to similarity
                )
                for r in rows
            ]

            # If vec_index returned 0 results, the index may be empty (data
            # only lives in embedding_index).  Fall back to brute-force so
            # retrieval still works while the index catches up.
            if not results:
                import logging as _logging
                _logging.getLogger(__name__).info(
                    "vec_index KNN returned 0 results; falling back to brute-force"
                )
                return await self._vec_search_brute_force(
                    conn, query_embedding, item_type, top_k, owner_id
                )

            return results
        except Exception:
            # vec_index may not be created yet; cache failure so subsequent
            # calls skip the attempt entirely instead of re-throwing each time.
            self._vec_loaded = False
            import logging as _logging
            _logging.getLogger(__name__).info("vec_index unavailable, falling back to brute-force search")
            return await self._vec_search_brute_force(
                conn, query_embedding, item_type, top_k, owner_id
            )

    #: Maximum number of embeddings to scan in brute-force mode before
    #: falling back to FTS-only search to prevent OOM on large datasets.
    MAX_BRUTE_FORCE_EMBEDDINGS: int = 10_000

    async def _vec_search_brute_force(
        self,
        conn: aiosqlite.Connection,
        query_embedding: List[float],
        item_type: Optional[str],
        top_k: int,
        owner_id: Optional[str] = None,
    ) -> List[VectorResult]:
        """Brute-force cosine similarity over embedding_index table.

        If the number of stored embeddings exceeds MAX_BRUTE_FORCE_EMBEDDINGS,
        logs a warning and returns an empty list so the caller can fall back
        to FTS-only search instead of loading everything into memory.

        Backward compatibility: ``owner_id`` is optional for legacy internal
        call sites/tests that invoked this helper before owner-aware filtering
        was added. When omitted, claim rows are not owner-filtered.
        """
        import logging
        _log = logging.getLogger(__name__)

        # Check count first to avoid OOM
        count_sql = "SELECT COUNT(*) FROM embedding_index"
        count_params: list = []
        if item_type:
            count_sql += " WHERE item_type = ?"
            count_params.append(item_type)
            if item_type == "claim" and owner_id is not None:
                count_sql += (
                    " AND item_id IN ("
                    "SELECT id FROM semantic_claim"
                    " WHERE owner_id = ? AND status = 'active'"
                    " AND (is_deleted = 0 OR is_deleted IS NULL))"
                )
                count_params.append(owner_id)
        else:
            if owner_id is not None:
                count_sql += (
                    " WHERE item_type = 'claim'"
                    " AND item_id IN ("
                    "SELECT id FROM semantic_claim"
                    " WHERE owner_id = ? AND status = 'active'"
                    " AND (is_deleted = 0 OR is_deleted IS NULL))"
                )
                count_params.append(owner_id)
            else:
                count_sql += " WHERE item_type = 'claim'"
        async with conn.execute(count_sql, count_params) as cur:
            row = await cur.fetchone()
        total = row[0] if row else 0

        if total > self.MAX_BRUTE_FORCE_EMBEDDINGS:
            _log.warning(
                "Brute-force vector search: %d embeddings exceed cap of %d. "
                "Returning empty results — caller should fall back to FTS.",
                total,
                self.MAX_BRUTE_FORCE_EMBEDDINGS,
            )
            return []

        if item_type:
            if item_type == "claim" and owner_id is not None:
                async with conn.execute(
                    """
                    SELECT ei.item_type, ei.item_id, ei.embedding, ei.model_id
                    FROM embedding_index ei
                    JOIN semantic_claim c ON c.id = ei.item_id
                    WHERE ei.item_type = ? AND c.owner_id = ?
                      AND c.status = 'active'
                      AND (c.is_deleted = 0 OR c.is_deleted IS NULL)
                    """,
                    (item_type, owner_id),
                ) as cur:
                    rows = await cur.fetchall()
            else:
                async with conn.execute(
                    "SELECT item_type, item_id, embedding, model_id FROM embedding_index WHERE item_type = ?",
                    (item_type,),
                ) as cur:
                    rows = await cur.fetchall()
        else:
            if owner_id is not None:
                async with conn.execute(
                    """
                    SELECT ei.item_type, ei.item_id, ei.embedding, ei.model_id
                    FROM embedding_index ei
                    JOIN semantic_claim c ON c.id = ei.item_id
                    WHERE ei.item_type = 'claim' AND c.owner_id = ?
                      AND c.status = 'active'
                      AND (c.is_deleted = 0 OR c.is_deleted IS NULL)
                    """,
                    (owner_id,),
                ) as cur:
                    rows = await cur.fetchall()
            else:
                async with conn.execute(
                    """
                    SELECT item_type, item_id, embedding, model_id
                    FROM embedding_index
                    WHERE item_type = 'claim'
                    """
                ) as cur:
                    rows = await cur.fetchall()

        scored: List[VectorResult] = []
        for row in rows:
            blob = row["embedding"]
            if not blob:
                continue
            vec = _unpack_floats(bytes(blob))
            if len(vec) != len(query_embedding):
                continue
            sim = _cosine_similarity(query_embedding, vec)
            scored.append(
                VectorResult(
                    item_type=row["item_type"],
                    item_id=row["item_id"],
                    score=sim,
                    model_id=row["model_id"],
                )
            )

        scored.sort(key=lambda r: r.score, reverse=True)
        return scored[:top_k]

    # ------------------------------------------------------------------ #
    # FTS operations
    # ------------------------------------------------------------------ #

    @staticmethod
    def _sanitize_fts_query(query: str) -> str:
        """Escape FTS5 special characters and join tokens with OR (#1575).

        FTS5 MATCH accepts operators (AND, OR, NOT, NEAR, *, ^) that could
        alter query semantics or cause syntax errors (#219).  We strip
        non-word/non-space characters so only plain tokens remain, then
        join with OR so multi-word queries match claims containing ANY
        token (BM25 still ranks multi-match higher).

        Previously tokens were joined with implicit AND, which killed
        BM25 candidates for queries like "Bicentennial Man connection"
        where not all words appear in any single claim.
        """
        # Remove FTS5 operators and special chars that break MATCH syntax
        sanitized = re.sub(r'[^\w\s]', ' ', query)
        tokens = sanitized.split()
        if not tokens:
            return '""'
        # Quote each token to prevent operator injection, join with OR
        # so any matching token produces BM25 candidates (#1575).
        return ' OR '.join(f'"{t}"' for t in tokens)

    @staticmethod
    def _sanitize_fts_query_strict(query: str) -> str:
        """Sanitize user-facing FTS5 query as a quoted literal phrase (#219).

        Wraps each token in double quotes so FTS5 operators embedded in
        user input (AND, OR, NOT, NEAR, *, ^) are treated as literals.
        """
        sanitized = re.sub(r'[^\w\s]', ' ', query)
        tokens = sanitized.split()
        if not tokens:
            return '""'
        # Quote each token individually so multi-word queries still match
        # documents containing all tokens (implicit AND), while preventing
        # operator injection.
        return ' '.join(f'"{t}"' for t in tokens)

    async def fts_search(
        self,
        query: str,
        table: str = "claims",
        top_k: int = 50,
        owner_id: str = "default",
        entity_id: Optional[str] = None,
        entity_type: Optional[str] = None,
        entity_ids: Optional[List[str]] = None,
    ) -> List[FTSResult]:
        conn = await self._get_conn()
        query = self._sanitize_fts_query(query)

        if table == "claims":
            # Build entity scoping clauses (#923)
            entity_clauses = ""
            entity_params: list = []
            if entity_id is not None:
                entity_clauses += " AND c.entity_id = ?"
                entity_params.append(entity_id)
            if entity_type is not None:
                entity_clauses += " AND c.entity_type = ?"
                entity_params.append(entity_type)
            # Support list-based entity scoping (#925)
            if entity_ids:
                placeholders = ", ".join("?" for _ in entity_ids)
                entity_clauses += (
                    f" AND (c.subject_entity_id IN ({placeholders})"
                    f" OR c.object_entity_id IN ({placeholders}))"
                )
                entity_params.extend(entity_ids)  # for subject_entity_id IN
                entity_params.extend(entity_ids)  # for object_entity_id IN
            async with conn.execute(
                f"""
                SELECT f.claim_id AS item_id,
                       f.content,
                       bm25(fts_claims) AS rank
                FROM fts_claims f
                JOIN semantic_claim c ON c.id = f.claim_id
                WHERE fts_claims MATCH ?
                  AND c.owner_id = ?
                  AND c.status = 'active'
                  AND (c.is_deleted = 0 OR c.is_deleted IS NULL)
                  {entity_clauses}
                ORDER BY rank
                LIMIT ?
                """,
                (query, owner_id, *entity_params, top_k),
            ) as cur:
                rows = await cur.fetchall()
        elif table == "episodes":
            async with conn.execute(
                """
                SELECT f.event_id AS item_id,
                       f.content,
                       bm25(fts_episodes) AS rank
                FROM fts_episodes f
                JOIN episode_event e ON e.id = f.event_id
                WHERE fts_episodes MATCH ?
                  AND e.owner_id = ?
                ORDER BY rank
                LIMIT ?
                """,
                (query, owner_id, top_k),
            ) as cur:
                rows = await cur.fetchall()
        else:
            raise ValueError(f"Unknown FTS table: {table!r}. Use 'claims' or 'episodes'.")

        return [
            FTSResult(item_id=r["item_id"], content=r["content"], rank=r["rank"])
            for r in rows
        ]
