"""可选的 API 嵌入，与稳定的 hash/sqlite-vec 索引相互隔离。

运行时的 model_id 是权威空间标识，涵盖提供商 URL、端点、模型与维度；维度相同并不表示兼容。
持久化向量用于按需构建各空间和维度的 sqlite-vec 索引。原生精确 KNN 避免每次搜索都由 Python
解码 JSON 并计算点积。覆盖率检查与排序使用同一事务。
"""

from __future__ import annotations

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

from app.database.db import connect_knowledge as connect, transaction
from app.errors import ApiError
from app.operation_logs import log_event
from app.retrieval.vectorstore import VectorHit
from app.retrieval.provenance import record_embedding
from app.retrieval.hybrid import rrf_fuse
from app.retrieval import space_index

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], *, local_only=False) -> EmbeddingResult: ...


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


def get_model_routing() -> EmbeddingRuntime | None:
    """惰性集成钩子；测试可以注入运行时而无需任何网络 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")
    # 缩放首先避免有限但极端的 API 值的上溢/下溢。
    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], *, accept_local=False, strict=False, local_only=False) -> RemoteEmbeddings | None:
    """返回经过验证的 API 向量，或 None 以使用调用者的本地基线。不要使用运行时的本地结果：调用者可能已经注入了自己的嵌入/存储对。异常特意排除取消。"""
    if not texts:
        return None
    try:
        runtime = get_model_routing()
        if runtime is None:
            if strict:
                raise ApiError(503, "EMBEDDING_UNAVAILABLE", "Embedding 服务未就绪，请检查模型路由和本地运行环境。")
            return None
        result = await runtime.embed(texts, local_only=True) if local_only else await runtime.embed(texts)
        if result.source != "api" and not accept_local:
            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],
            source=result.source,
        )
    except Exception as exc:
        log_event('vectors', 'embedding.failed', level='ERROR' if strict else 'WARNING', error=exc,
                  count=len(texts), fallback='none' if strict else 'local_index')
        # 避免记录包含凭据或笔记文本的提供程序异常。
        record_embedding(fallback_reason="REMOTE_EMBEDDING_UNAVAILABLE")
        logger.warning("Remote embedding unavailable (%s); using local index", type(exc).__name__)
        if strict:
            if isinstance(exc, ApiError):
                raise
            raise ApiError(503, "EMBEDDING_UNAVAILABLE", "Embedding 调用失败或返回无效，请检查模型路由、API 和本地模型运行状态。") from exc
        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, dimensions, block_id)
        )
    """)
    primary = [row[1] for row in sorted(conn.execute('PRAGMA table_info(routed_block_vectors)'), key=lambda row: row[5]) if row[5]]
    if primary == ['space_id', 'block_id']:
        conn.execute('CREATE TABLE routed_block_vectors_upgrade (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,dimensions,block_id))')
        conn.execute('INSERT INTO routed_block_vectors_upgrade SELECT * FROM routed_block_vectors')
        conn.execute('DROP TABLE routed_block_vectors')
        conn.execute('ALTER TABLE routed_block_vectors_upgrade RENAME TO routed_block_vectors')
    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:
    """在调用方的元数据事务内尽力写入辅助索引。

    savepoint 可阻止只写入部分远程批次，并将存储故障与笔记保存隔离；替换或删除内容块时，
    所有旧空间都会自动级联清理。
    """
    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, dimensions, 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)
                ],
            )
            space_index.upsert(conn, block_ids, batch)
        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, accept_local=False, strict=False) -> list[VectorHit] | None:
    """None 表示回退，包括任何丢失/无效的当前块向量。将覆盖率和向量一起读取，以便并发笔记更新无法生成明显完整的子集。切勿用本地命中来填补缺失的远程命中。"""
    if accept_local:
        conn = connect()
        try:
            policies = {bool(row[0]) for row in conn.execute("SELECT DISTINCT embedding_local_only FROM blocks")}
        finally:
            conn.close()
        if True in policies:
            return await _search_partitioned(query, policies, top_k=top_k, strict=strict)
    batch = await embed_remote([query], accept_local=accept_local, strict=strict)
    if batch is None:
        return None

    if not await _prepare_for_search([batch], strict):
        return None
    return await asyncio.to_thread(_search_space, batch, top_k, strict)


async def _prepare_indexes(batches):
    from app.services.coordination import vault_mutation_lock
    def prepare(check_only=False):
        conn = connect()
        try:
            if check_only:
                return space_index.is_ready(conn, batches)
            space_index.prepare(conn, batches)
        finally:
            conn.close()
    if await asyncio.to_thread(prepare, True):
        return
    # 与保存共享协作门：当迁移在另一个线程中拥有 SQLite 写锁时，永远不会阻塞 SQLite 写锁上的事件循环。
    async with vault_mutation_lock():
        work = asyncio.create_task(asyncio.to_thread(prepare))
        cancelled = False
        while not work.done():
            try:
                await asyncio.shield(work)
            except asyncio.CancelledError:
                cancelled = True
        work.result()
        if cancelled:
            raise asyncio.CancelledError


async def _prepare_for_search(batches, strict):
    try:
        await _prepare_indexes(batches)
        return True
    except Exception as exc:
        record_embedding(fallback_reason='REMOTE_INDEX_UNAVAILABLE')
        if strict:
            raise ApiError(409, 'SEMANTIC_INDEX_UNAVAILABLE', '向量索引准备失败，请检查索引状态。') from exc
        return False


def _search_space(batch, top_k, strict):
    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")
                    if not conn.execute("SELECT 1 FROM blocks LIMIT 1").fetchone():
                        return []
                    if strict:
                        raise ValueError("semantic index missing")
                    return None
                result = space_index.search(conn, batch, top_k)
                record_embedding(source=batch.source, 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__)
        if strict:
            raise ApiError(409, "SEMANTIC_INDEX_UNAVAILABLE",
                           "Embedding 已可用，但当前模型的向量索引缺失、不完整或已失效。请在「设置 → 索引与模型」中重建全部索引。",
                           {"model_id": batch.space_id, "dimensions": batch.dimensions, "source": batch.source}) from exc
        return None


async def _search_partitioned(query: str, policies: set[bool], *, top_k: int, strict: bool):
    """按策略嵌入；独立对每个空间进行排名并融合排名，而不是向量。"""
    batches = {}
    for policy in sorted(policies):
        batch = await embed_remote([query], accept_local=True, strict=strict, local_only=policy)
        if batch is None:
            return None
        batches[policy] = batch
    if not await _prepare_for_search(list(batches.values()), strict):
        return None
    return await asyncio.to_thread(_search_partitions, batches, policies, top_k, strict)


def _search_partitions(batches, policies, top_k, strict):
    conn = connect()
    try:
        with transaction(conn):
            # 在打开单个读取快照之前，查询向量已准备就绪。
            current = {bool(row[0]) for row in conn.execute("SELECT DISTINCT embedding_local_only FROM blocks")}
            if current != policies:
                raise ValueError("embedding policies changed while querying")
            ranked = []
            for policy, batch in batches.items():
                ranked.append(space_index.search(conn, batch, top_k, policy))
            spaces = [{"source": b.source, "model_id": b.space_id, "dimensions": b.dimensions,
                       "local_only": policy} for policy, b in batches.items()]
            record_embedding(source="mixed" if len({b.source for b in batches.values()}) > 1 else batch.source,
                             spaces=spaces, fallback_reason=None)
            if len(ranked) == 1:
                return ranked[0]
            fused = rrf_fuse([[hit.id for hit in group] for group in ranked])
            return [VectorHit(id=key, score=score) for key, score in
                    sorted(fused.items(), key=lambda item: (-item[1], item[0]))[:top_k]]
    except Exception as exc:
        record_embedding(source="unavailable", fallback_reason="REMOTE_INDEX_UNAVAILABLE")
        if strict:
            raise ApiError(409, "SEMANTIC_INDEX_UNAVAILABLE", "部分索引分区缺失或已失效，请重建全部索引。") from exc
        return None
    finally:
        conn.close()