"""Embedding function helpers and cache support."""

from __future__ import annotations

import json
import logging
import sqlite3
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Mapping, Optional, Sequence

from .index_state import sha256_text

logger = logging.getLogger(__name__)


class HashEmbeddingFunction:
    """Deterministic fallback embedding without external model dependencies."""

    def __init__(self, model_name: str = "hash-embedding-v1", dimension: int = 64) -> None:
        self.model_name = model_name
        self.dimension = max(16, int(dimension))

    def _embed_one(self, text: str) -> list[float]:
        digest = sha256_text(text)
        seed = bytes.fromhex(digest)
        raw = (seed * ((self.dimension // len(seed)) + 1))[: self.dimension]
        # map bytes to [-1, 1]
        return [round((byte / 255.0) * 2.0 - 1.0, 6) for byte in raw]

    def __call__(self, input: Sequence[str]) -> list[list[float]]:  # noqa: A002 - chroma API
        return [self._embed_one(text or "") for text in input]

    def embed_documents(
        self,
        input: Optional[Sequence[str]] = None,  # noqa: A002 - compatibility
        **kwargs,
    ) -> list[list[float]]:
        texts = input if input is not None else kwargs.get("texts") or []
        if isinstance(texts, str):
            texts = [texts]
        return self.__call__(texts)

    def embed_query(
        self,
        input: Optional[str] = None,  # noqa: A002 - compatibility
        **kwargs,
    ) -> list[list[float]]:
        text = input if input is not None else kwargs.get("text") or ""
        if isinstance(text, (list, tuple)):
            texts = [str(item) for item in text]
        else:
            texts = [str(text)]
        return self.__call__(texts)

    def name(self) -> str:
        return self.model_name

    def is_legacy(self) -> bool:
        return True


class EmbeddingCache:
    """SQLite embedding cache keyed by (model_id, chunk_hash)."""

    def __init__(self, cache_path: str | Path) -> None:
        self.path = Path(cache_path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(str(self.path))
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS embedding_cache (
              model_id TEXT NOT NULL,
              chunk_hash TEXT NOT NULL,
              vector_json TEXT NOT NULL,
              updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
              PRIMARY KEY (model_id, chunk_hash)
            )
            """
        )
        self._conn.commit()

    def fetch(self, model_id: str, chunk_hashes: Iterable[str]) -> dict[str, list[float]]:
        hashes = [item for item in chunk_hashes]
        if not hashes:
            return {}
        placeholders = ",".join("?" for _ in hashes)
        cursor = self._conn.execute(
            f"""
            SELECT chunk_hash, vector_json
            FROM embedding_cache
            WHERE model_id = ? AND chunk_hash IN ({placeholders})
            """,
            [model_id, *hashes],
        )
        rows = cursor.fetchall()
        result: dict[str, list[float]] = {}
        for chunk_hash, vector_json in rows:
            try:
                value = json.loads(vector_json)
                if isinstance(value, list):
                    result[str(chunk_hash)] = [float(v) for v in value]
            except Exception:
                continue
        return result

    def upsert(self, model_id: str, vectors: Mapping[str, Sequence[float]]) -> None:
        rows = [
            (model_id, chunk_hash, json.dumps([float(v) for v in vector]))
            for chunk_hash, vector in vectors.items()
        ]
        if not rows:
            return
        self._conn.executemany(
            """
            INSERT INTO embedding_cache(model_id, chunk_hash, vector_json)
            VALUES (?, ?, ?)
            ON CONFLICT(model_id, chunk_hash) DO UPDATE SET
              vector_json = excluded.vector_json,
              updated_at = CURRENT_TIMESTAMP
            """,
            rows,
        )
        self._conn.commit()

    def close(self) -> None:
        self._conn.close()


class CachedEmbeddingFunction:
    """Cache wrapper for any chroma-compatible embedding function."""

    def __init__(self, base_function: Any, model_id: str, cache: EmbeddingCache) -> None:
        self.base_function = base_function
        self.model_id = model_id
        self.cache = cache

    def __call__(self, input: Sequence[str]) -> list[list[float]]:  # noqa: A002 - chroma API
        texts = [text or "" for text in input]
        hashes = [sha256_text(text) for text in texts]
        cache_hits = self.cache.fetch(self.model_id, hashes)

        result: list[Optional[list[float]]] = [None] * len(texts)
        missing_indices: list[int] = []
        missing_texts: list[str] = []
        missing_hashes: list[str] = []

        for idx, chunk_hash in enumerate(hashes):
            cached = cache_hits.get(chunk_hash)
            if cached is not None:
                result[idx] = cached
            else:
                missing_indices.append(idx)
                missing_texts.append(texts[idx])
                missing_hashes.append(chunk_hash)

        if missing_texts:
            generated = self.base_function(missing_texts)
            generated_vectors: dict[str, Sequence[float]] = {}
            for local_idx, vector in enumerate(generated):
                index = missing_indices[local_idx]
                final_vector = [float(v) for v in vector]
                result[index] = final_vector
                generated_vectors[missing_hashes[local_idx]] = final_vector
            self.cache.upsert(self.model_id, generated_vectors)

        return [vector or [] for vector in result]

    def embed_documents(
        self,
        input: Optional[Sequence[str]] = None,  # noqa: A002 - compatibility
        **kwargs,
    ) -> list[list[float]]:
        texts = input if input is not None else kwargs.get("texts") or []
        if isinstance(texts, str):
            texts = [texts]
        return self.__call__(texts)

    def embed_query(
        self,
        input: Optional[str] = None,  # noqa: A002 - compatibility
        **kwargs,
    ) -> list[list[float]]:
        text = input if input is not None else kwargs.get("text") or ""
        if isinstance(text, (list, tuple)):
            texts = [str(item) for item in text]
        else:
            texts = [str(text)]
        return self.__call__(texts)

    def name(self) -> str:
        if hasattr(self.base_function, "name"):
            return str(self.base_function.name())
        return self.model_id

    def is_legacy(self) -> bool:
        if hasattr(self.base_function, "is_legacy"):
            return bool(self.base_function.is_legacy())
        return True


@dataclass
class EmbeddingRuntime:
    model_id: str
    embedding_function: Any
    cache: Optional[EmbeddingCache] = None

    def close(self) -> None:
        if self.cache is not None:
            self.cache.close()


def _build_sentence_transformer_embedding(model_name: str):
    from chromadb.utils import embedding_functions

    return embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name)


def resolve_embedding_runtime(
    embedding_model: Optional[str] = None,
    *,
    cache_path: Optional[str | Path] = None,
) -> EmbeddingRuntime:
    """
    Resolve embedding function with graceful fallback.

    If sentence-transformers is unavailable, falls back to deterministic hash embeddings.
    """
    selected_model = (embedding_model or "").strip()
    cache = EmbeddingCache(cache_path) if cache_path else None

    if not selected_model:
        base = HashEmbeddingFunction()
        function = CachedEmbeddingFunction(base, base.model_name, cache) if cache else base
        return EmbeddingRuntime(model_id=base.model_name, embedding_function=function, cache=cache)

    model_name = selected_model
    model_path = Path(selected_model)
    if model_path.exists():
        model_name = str(model_path.resolve())

    try:
        base = _build_sentence_transformer_embedding(model_name)
        if cache:
            wrapped = CachedEmbeddingFunction(base, model_name, cache)
            return EmbeddingRuntime(model_id=model_name, embedding_function=wrapped, cache=cache)
        return EmbeddingRuntime(model_id=model_name, embedding_function=base, cache=None)
    except Exception as exc:
        logger.warning("SentenceTransformer 嵌入不可用，回退 hash 嵌入: %s", exc)
        fallback = HashEmbeddingFunction(model_name=f"hash::{model_name}")
        wrapped = CachedEmbeddingFunction(fallback, fallback.model_name, cache) if cache else fallback
        return EmbeddingRuntime(model_id=fallback.model_name, embedding_function=wrapped, cache=cache)


__all__ = [
    "EmbeddingCache",
    "EmbeddingRuntime",
    "CachedEmbeddingFunction",
    "HashEmbeddingFunction",
    "resolve_embedding_runtime",
]
