"""Optional API embeddings, isolated from the stable hash/sqlite-vec index.

The runtime's model_id is the authoritative space ID (including provider URL,
endpoint, model and dimensions); equal dimensions alone never imply compatibility.
This phase uses a lazy, rebuildable SQLite side table instead of a schema migration.
Search scans only current blocks in one database snapshot and requires complete
coverage. Cosine ranking costs O(blocks * dimensions) with an O(top_k) heap; this
small-vault implementation should become a per-space ANN index at larger scale.
"""

from __future__ import annotations

import heapq
import json
import logging
import math
import sqlite3
from dataclasses import dataclass
from typing import Protocol

from app.database.db import connect, transaction
from app.retrieval.vectorstore import VectorHit
from app.retrieval.provenance import record_embedding

logger = logging.getLogger(__name__)


class EmbeddingResult(Protocol):
    vectors: list[list[float]]
    source: str
    model_id: str
    dimensions: int
    fallback_reason: str | None


class EmbeddingRuntime(Protocol):
    async def embed(self, texts: list[str]) -> EmbeddingResult: ...


@dataclass(frozen=True)
class RemoteEmbeddings:
    space_id: str
    dimensions: int
    vectors: list[list[float]]


def get_model_routing() -> EmbeddingRuntime | None:
    """Lazy integration hook; tests can inject a runtime without any network I/O."""
    from app.container import container

    return getattr(container, "model_routing", None)


def _unit_vector(vector: list[float], dimensions: int) -> list[float]:
    if len(vector) != dimensions:
        raise ValueError("embedding dimension mismatch")
    if any(isinstance(value, bool) or not isinstance(value, (int, float)) for value in vector):
        raise ValueError("embedding must be numeric")
    if not all(math.isfinite(value) for value in vector):
        raise ValueError("embedding must be finite")
    scale = max(abs(value) for value in vector)
    if scale == 0:
        raise ValueError("embedding must be nonzero")
    # Scaling first avoids overflow/underflow for finite but extreme API values.
    scaled = [value / scale for value in vector]
    norm = math.sqrt(math.fsum(value * value for value in scaled))
    return [value / norm for value in scaled]


async def embed_remote(texts: list[str]) -> RemoteEmbeddings | None:
    """Return validated API vectors, or None to use the caller's local baseline.

    Do not use the runtime's local result: the caller may have injected its own
    embedding/store pair. Exception deliberately excludes cancellation.
    """
    if not texts:
        return None
    try:
        runtime = get_model_routing()
        if runtime is None:
            return None
        result = await runtime.embed(texts)
        if result.source != "api":
            record_embedding(fallback_reason=result.fallback_reason)
            return None
        if not isinstance(result.model_id, str) or not result.model_id or result.model_id == "hash-v1":
            raise ValueError("API embedding needs a distinct space ID")
        if type(result.dimensions) is not int or result.dimensions <= 0:
            raise ValueError("invalid embedding dimensions")
        if len(result.vectors) != len(texts):
            raise ValueError("embedding count mismatch")
        return RemoteEmbeddings(
            space_id=result.model_id,
            dimensions=result.dimensions,
            vectors=[_unit_vector(vector, result.dimensions) for vector in result.vectors],
        )
    except Exception as exc:
        # Avoid logging provider exceptions containing credentials or note text.
        record_embedding(fallback_reason="REMOTE_EMBEDDING_UNAVAILABLE")
        logger.warning("Remote embedding unavailable (%s); using local index", type(exc).__name__)
        return None


def _ensure_table(conn: sqlite3.Connection) -> None:
    conn.execute("""
        CREATE TABLE IF NOT EXISTS routed_block_vectors (
            space_id TEXT NOT NULL,
            block_id TEXT NOT NULL REFERENCES blocks(block_id) ON DELETE CASCADE,
            dimensions INTEGER NOT NULL CHECK (dimensions > 0),
            vector TEXT NOT NULL,
            PRIMARY KEY (space_id, block_id)
        )
    """)
    conn.execute("""
        CREATE INDEX IF NOT EXISTS routed_block_vectors_block_id
        ON routed_block_vectors(block_id)
    """)


def store_remote(
    conn: sqlite3.Connection, block_ids: list[str], batch: RemoteEmbeddings | None,
) -> None:
    """Best-effort side-index write inside the caller's metadata transaction.

    A savepoint prevents partial remote batches and isolates storage failures from
    note saving. Replacing/deleting blocks cascades all old spaces automatically.
    """
    if batch is None:
        return
    try:
        conn.execute("SAVEPOINT routed_vectors_write")
        try:
            if len(block_ids) != len(batch.vectors):
                raise ValueError("block/vector count mismatch")
            _ensure_table(conn)
            conn.executemany(
                """INSERT INTO routed_block_vectors (space_id, block_id, dimensions, vector)
                   VALUES (?, ?, ?, ?)
                   ON CONFLICT (space_id, block_id) DO UPDATE SET
                       dimensions = excluded.dimensions, vector = excluded.vector""",
                [
                    (batch.space_id, block_id, batch.dimensions, json.dumps(vector, allow_nan=False))
                    for block_id, vector in zip(block_ids, batch.vectors)
                ],
            )
        except BaseException:
            conn.execute("ROLLBACK TO routed_vectors_write")
            raise
        finally:
            conn.execute("RELEASE routed_vectors_write")
    except Exception as exc:
        logger.warning("Remote vector storage unavailable (%s); local index retained", type(exc).__name__)


async def search_remote(query: str, *, top_k: int) -> list[VectorHit] | None:
    """None means fallback, including any missing/invalid current-block vector.

    Read coverage and vectors together so concurrent note updates cannot produce
    an apparently complete subset. Never fill missing remote hits with local hits.
    """
    batch = await embed_remote([query])
    if batch is None:
        return None
    record_embedding(attempted_space={"model_id": batch.space_id, "dimensions": batch.dimensions})
    try:
        conn = connect()
        try:
            with transaction(conn):
                exists = conn.execute(
                    "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'routed_block_vectors'"
                ).fetchone()
                if exists is None:
                    record_embedding(fallback_reason="REMOTE_INDEX_MISSING")
                    return None
                rows = conn.execute(
                    """SELECT b.block_id, r.vector
                       FROM blocks AS b
                       LEFT JOIN routed_block_vectors AS r
                         ON r.block_id = b.block_id AND r.space_id = ? AND r.dimensions = ?
                       ORDER BY b.block_id""",
                    (batch.space_id, batch.dimensions),
                )

                def hits():
                    for row in rows:
                        if row["vector"] is None:
                            raise ValueError("remote space has incomplete block coverage")
                        vector = _unit_vector(json.loads(row["vector"]), batch.dimensions)
                        score = math.fsum(a * b for a, b in zip(batch.vectors[0], vector))
                        yield VectorHit(id=row["block_id"], score=max(0.0, min(1.0, score)))

                result = heapq.nlargest(top_k, hits(), key=lambda hit: hit.score)
                record_embedding(source="api", model_id=batch.space_id,
                                 dimensions=batch.dimensions, fallback_reason=None)
                return result
        finally:
            conn.close()
    except Exception as exc:
        record_embedding(fallback_reason="REMOTE_INDEX_UNAVAILABLE")
        logger.debug("Remote vector search unavailable (%s); using local index", type(exc).__name__)
        return None