"""Claim search operations mixin (entity, status, timeline, provenance, embeddings)."""
from __future__ import annotations

from typing import List, Optional

from cortex.storage._helpers import (
    _now, _pack_floats, _unpack_floats, _row_to_claim, _row_to_provenance,
)
from cortex.types import ClaimProvenance, SemanticClaim, VectorResult


class ClaimSearchMixin:
    """Mixin providing claim search and embedding storage."""

    async def get_claims_by_entity(
        self,
        entity_id: str,
        owner_id: Optional[str] = None,
    ) -> List[SemanticClaim]:
        """Return active claims linked to a resolved entity_id.

        Notes:
            - API-facing call sites must pass ``owner_id`` to enforce tenant
              isolation.
            - ``owner_id=None`` is retained for internal-only compatibility.
        """
        conn = await self._get_conn()
        if owner_id is None:
            async with conn.execute(
                """
                SELECT * FROM semantic_claim
                WHERE subject_entity_id = ? AND status = 'active'
                  AND (is_deleted = 0 OR is_deleted IS NULL)
                ORDER BY confidence DESC
                """,
                (entity_id,),
            ) as cur:
                rows = await cur.fetchall()
        else:
            async with conn.execute(
                """
                SELECT * FROM semantic_claim
                WHERE owner_id = ? AND subject_entity_id = ? AND status = 'active'
                  AND (is_deleted = 0 OR is_deleted IS NULL)
                ORDER BY confidence DESC
                """,
                (owner_id, entity_id),
            ) as cur:
                rows = await cur.fetchall()
        return [_row_to_claim(r) for r in rows]

    async def get_claims_by_status(
        self, status: str, owner_id: str = "default"
    ) -> List[SemanticClaim]:
        """Return all claims with a given status for an owner."""
        conn = await self._get_conn()
        async with conn.execute(
            """
            SELECT * FROM semantic_claim
            WHERE owner_id = ? AND status = ?
              AND (is_deleted = 0 OR is_deleted IS NULL)
            ORDER BY updated_at DESC
            """,
            (owner_id, status),
        ) as cur:
            rows = await cur.fetchall()
        return [_row_to_claim(r) for r in rows]

    async def get_claims_timeline(
        self, subject: str, owner_id: str = "default"
    ) -> List["SemanticClaim"]:
        """Return ALL claims (all statuses) for a subject, ordered by valid_from (#925)."""
        conn = await self._get_conn()
        async with conn.execute(
            """
            SELECT * FROM semantic_claim
            WHERE owner_id = ? AND subject = ?
              AND (is_deleted = 0 OR is_deleted IS NULL)
            ORDER BY COALESCE(valid_from, created_at) ASC
            """,
            (owner_id, subject),
        ) as cur:
            rows = await cur.fetchall()
        return [_row_to_claim(r) for r in rows]

    async def add_provenance(
        self, claim_id: str, episode_event_id: str, **kwargs
    ) -> ClaimProvenance:
        async with self._write_lock:
            conn = await self._get_conn()
            now = _now()
            async with conn.execute(
                """
                INSERT INTO claim_provenance
                    (claim_id, episode_event_id, span_start, span_end,
                     extraction_method, extractor_model_id,
                     extraction_prompt_ver, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                RETURNING *
                """,
                (
                    claim_id,
                    episode_event_id,
                    kwargs.get("span_start"),
                    kwargs.get("span_end"),
                    kwargs.get("extraction_method"),
                    kwargs.get("extractor_model_id"),
                    kwargs.get("extraction_prompt_ver"),
                    now,
                ),
            ) as cur:
                row = await cur.fetchone()
            await conn.commit()
            return _row_to_provenance(row)

    async def get_provenance(self, claim_id: str) -> list:
        """Return all provenance records for a given claim."""
        conn = await self._get_conn()
        async with conn.execute(
            "SELECT * FROM claim_provenance WHERE claim_id = ? ORDER BY created_at ASC",
            (claim_id,),
        ) as cur:
            rows = await cur.fetchall()
        return [_row_to_provenance(r) for r in rows]

    async def store_embedding(
        self,
        item_type: str,
        item_id: str,
        embedding: List[float],
        model_id: str,
    ) -> None:
        import math
        import logging
        _log = logging.getLogger(__name__)

        # Validate embedding vector before storage
        if not embedding:
            _log.warning("store_embedding: empty embedding for %s/%s, skipping", item_type, item_id)
            return
        if any(math.isnan(x) or math.isinf(x) for x in embedding):
            _log.warning("store_embedding: NaN/Inf in embedding for %s/%s, skipping", item_type, item_id)
            return
        if all(x == 0.0 for x in embedding):
            _log.warning("store_embedding: zero-vector embedding for %s/%s, skipping", item_type, item_id)
            return

        async with self._write_lock:
            conn = await self._get_conn()
            now = _now()
            blob = _pack_floats(embedding)

            # Upsert: delete old, insert new
            await conn.execute(
                "DELETE FROM embedding_index WHERE item_type = ? AND item_id = ? AND model_id = ?",
                (item_type, item_id, model_id),
            )
            await conn.execute(
                """
                INSERT INTO embedding_index (item_type, item_id, embedding, model_id, created_at)
                VALUES (?, ?, ?, ?, ?)
                """,
                (item_type, item_id, blob, model_id, now),
            )

            # Also insert into vec_index (sqlite-vec KNN virtual table) if available.
            # The vec0 virtual table uses rowid-based upsert, so we delete first.
            if self._vec_loaded:
                try:
                    await conn.execute(
                        "DELETE FROM vec_index WHERE item_type = ? AND item_id = ?",
                        (item_type, item_id),
                    )
                    await conn.execute(
                        "INSERT INTO vec_index (item_type, item_id, embedding) VALUES (?, ?, ?)",
                        (item_type, item_id, blob),
                    )
                except Exception as vec_exc:
                    import logging as _logging
                    _logging.getLogger(__name__).warning(
                        "vec_index insert failed for %s/%s — brute-force will be used: %s",
                        item_type, item_id, vec_exc,
                    )

            await conn.commit()

    async def vector_search(
        self,
        query_embedding: List[float],
        item_type: Optional[str] = None,
        top_k: int = 50,
        owner_id: str = "default",
    ) -> List[VectorResult]:
        conn = await self._get_conn()

        if self._vec_loaded:
            return await self._vec_search_sqlite_vec(
                conn, query_embedding, item_type, top_k, owner_id
            )
        return await self._vec_search_brute_force(
            conn, query_embedding, item_type, top_k, owner_id
        )
