from __future__ import annotations

import os
import sys
import threading
from abc import ABC, abstractmethod
from typing import Iterable, Sequence

import faiss
import numpy as np
import httpx

from mikoshi.config import Config, ConfigError


class EmbeddingsProvider(ABC):
    name: str
    model: str
    dimension: int

    @abstractmethod
    def embed_texts(self, texts: Sequence[str]) -> np.ndarray:
        raise NotImplementedError


def normalize_embeddings(vectors: np.ndarray) -> np.ndarray:
    if vectors.size:
        faiss.normalize_L2(vectors)
    return vectors


def _batch_items(items: Sequence[str], batch_size: int) -> Iterable[list[str]]:
    for i in range(0, len(items), batch_size):
        yield list(items[i : i + batch_size])


class LocalEmbeddingsProvider(EmbeddingsProvider):
    def __init__(self, model_name: str) -> None:
        os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
        from huggingface_hub import snapshot_download
        from sentence_transformers import SentenceTransformer
        offline = os.getenv("MIKOSHI_OFFLINE", "").strip().lower() in {
            "1",
            "true",
            "yes",
            "on",
        }
        timeout_sec = int(os.getenv("MIKOSHI_MODEL_DOWNLOAD_TIMEOUT_SEC", "300"))
        try:
            model_path = snapshot_download(model_name, local_files_only=True)
        except Exception as exc:  # pragma: no cover - depends on local cache
            if offline:
                raise ConfigError(
                    "Local embeddings model not found in cache. "
                    "Download it once with: python -c \"from huggingface_hub import snapshot_download; "
                    f"snapshot_download('{model_name}')\""
                ) from exc
            print("⬇️ Downloading local model (one-time)…", file=sys.stderr)
            result: dict[str, object] = {}

            def _download() -> None:
                try:
                    result["path"] = snapshot_download(
                        model_name, local_files_only=False
                    )
                except Exception as download_exc:  # pragma: no cover - depends on network
                    result["error"] = download_exc

            thread = threading.Thread(target=_download, daemon=True)
            thread.start()
            thread.join(timeout=timeout_sec)
            if thread.is_alive():
                raise ConfigError(
                    "Model download timed out. Set MIKOSHI_OFFLINE=1 to skip, or retry on a stable network."
                )
            if "error" in result:
                raise ConfigError(
                    "Failed to download local embeddings model."
                ) from result["error"]
            model_path = str(result["path"])
            print("✅ Model ready", file=sys.stderr)

        self.name = "local"
        self.model = model_name
        print("🧠 Loading embeddings model…", file=sys.stderr)
        self._model = SentenceTransformer(model_path, device="cpu")
        try:
            import torch

            torch.set_num_threads(1)
            torch.set_num_interop_threads(1)
        except Exception:
            pass
        print("✅ Embeddings model loaded", file=sys.stderr)
        self.dimension = int(self._model.get_sentence_embedding_dimension())

    def embed_texts(self, texts: Sequence[str]) -> np.ndarray:
        if not texts:
            return np.zeros((0, self.dimension), dtype=np.float32)
        embeddings = self._model.encode(
            list(texts),
            batch_size=32,
            show_progress_bar=False,
            convert_to_numpy=True,
        )
        return np.asarray(embeddings, dtype=np.float32)


class OpenAIEmbeddingsProvider(EmbeddingsProvider):
    def __init__(self, api_key: str, base_url: str | None, model: str) -> None:
        if not api_key:
            raise ConfigError(
                "OpenAI API key missing. Set MIKOSHI_OPENAI_API_KEY or OPENAI_API_KEY."
            )
        self.name = "openai"
        self.model = model
        self.dimension = 0
        self._api_key = api_key
        self._base_url = (base_url or "https://api.openai.com/v1").rstrip("/")
        self._client = httpx.Client(timeout=30.0)

    def _endpoint(self) -> str:
        return f"{self._base_url}/embeddings"

    def embed_texts(self, texts: Sequence[str]) -> np.ndarray:
        if not texts:
            return np.zeros((0, 0), dtype=np.float32)
        embeddings: list[list[float]] = []
        for batch in _batch_items(list(texts), 96):
            response = self._client.post(
                self._endpoint(),
                headers={"Authorization": f"Bearer {self._api_key}"},
                json={"model": self.model, "input": batch},
            )
            response.raise_for_status()
            payload = response.json()
            data = sorted(payload.get("data", []), key=lambda item: item.get("index", 0))
            embeddings.extend([item["embedding"] for item in data])

        array = np.asarray(embeddings, dtype=np.float32)
        if array.size and self.dimension == 0:
            self.dimension = int(array.shape[1])
        return array


def get_embeddings_provider(config: Config) -> EmbeddingsProvider:
    provider = config.embeddings.provider
    if provider == "local":
        return LocalEmbeddingsProvider(config.embeddings.model)
    if provider == "openai":
        return OpenAIEmbeddingsProvider(
            api_key=config.embeddings.openai_api_key or "",
            base_url=config.embeddings.openai_base_url,
            model=config.embeddings.model,
        )
    raise ConfigError(f"Unknown embeddings provider: {provider}")


class SemanticSearcher:
    def __init__(self, index: faiss.Index, provider: EmbeddingsProvider) -> None:
        self.index = index
        self.provider = provider

    def search(self, query: str, k: int) -> list[tuple[int, float]]:
        query_vec = self.provider.embed_texts([query])
        if query_vec.size == 0:
            return []
        query_vec = normalize_embeddings(query_vec)
        scores, indices = self.index.search(query_vec, k)
        results: list[tuple[int, float]] = []
        for score, idx in zip(scores[0], indices[0]):
            if idx < 0:
                continue
            results.append((int(idx), float(score)))
        return results
