"""Multi-model embedding router with dimension normalization."""

from __future__ import annotations

import json
import logging
import time
import urllib.request
from typing import Any

import numpy as np

from neurovault.cognitive.embedding_engine import (
    EmbeddingBackend,
    OllamaEmbeddingBackend,
    SentenceTransformerBackend,
    TFIDFBackend,
)

logger = logging.getLogger(__name__)


def _normalize_dimension(vector: list[float], target_dim: int) -> list[float]:
    arr = np.asarray(vector, dtype=np.float32)
    if arr.size >= target_dim:
        arr = arr[:target_dim]
    else:
        arr = np.pad(arr, (0, target_dim - arr.size), mode="constant")
    norm = float(np.linalg.norm(arr))
    if norm > 0:
        arr = arr / norm
    return arr.tolist()


class OpenAICompatBackend(EmbeddingBackend):
    """Embedding backend for OpenAI-compatible API endpoints."""

    name = "openai-compat"

    def __init__(
        self,
        base_url: str,
        api_key: str = "",
        model: str = "text-embedding-3-small",
        timeout_s: float = 10.0,
    ):
        self._base_url = str(base_url).rstrip("/")
        self._api_key = api_key
        self._model = model
        self._timeout_s = float(timeout_s)
        probe = self.embed("probe")
        self._dim = len(probe)
        if self._dim <= 0:
            raise ValueError("openai-compatible embedding backend returned empty vector")

    @property
    def dimension(self) -> int:
        return self._dim

    def embed(self, text: str) -> list[float]:
        payload = json.dumps(
            {
                "model": self._model,
                "input": text,
            }
        ).encode("utf-8")
        headers = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        req = urllib.request.Request(
            f"{self._base_url}/v1/embeddings",
            data=payload,
            headers=headers,
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=self._timeout_s) as resp:
            result = json.loads(resp.read().decode("utf-8"))

        data = result.get("data", [])
        if not data:
            raise ValueError("openai-compatible response missing data")
        embedding = data[0].get("embedding")
        if not isinstance(embedding, list):
            raise ValueError("openai-compatible response missing embedding vector")
        return [float(v) for v in embedding]


class PCAReducer:
    """Reduce embedding dimensions via PCA for cross-model compatibility."""

    def __init__(self, target_dim: int = 384):
        self._target_dim = int(target_dim)
        self._transform: np.ndarray | None = None
        self._mean: np.ndarray | None = None

    def fit(self, vectors: list[list[float]]) -> None:
        """Fit PCA transform from a sample of vectors."""
        if not vectors:
            return

        matrix = np.array(vectors, dtype=np.float32)
        if matrix.shape[1] <= self._target_dim:
            self._transform = None
            self._mean = None
            return

        mean = matrix.mean(axis=0)
        centered = matrix - mean
        _, _, vt = np.linalg.svd(centered, full_matrices=False)
        self._transform = vt[: self._target_dim].T
        self._mean = mean

    def reduce(self, vector: list[float]) -> list[float]:
        if len(vector) <= self._target_dim:
            return _normalize_dimension(vector, self._target_dim)
        if self._transform is None or self._mean is None:
            return _normalize_dimension(vector, self._target_dim)

        v = np.asarray(vector, dtype=np.float32) - self._mean
        reduced = v @ self._transform
        return _normalize_dimension(reduced.tolist(), self._target_dim)


class EmbeddingRouter:
    """Route embedding requests to the best available backend."""

    def __init__(self, config: dict[str, Any] | None = None):
        self._config = config or {}
        self._target_dim = int(self._config.get("target_dimension", 384))
        self._backends: list[EmbeddingBackend] = []
        self._active: EmbeddingBackend | None = None
        self._reducer: PCAReducer | None = None
        self._initialize()

    @property
    def dimension(self) -> int:
        return self._target_dim

    @property
    def active_backend(self) -> str:
        return self._active.name if self._active else "none"

    @property
    def available_backends(self) -> list[str]:
        return [backend.name for backend in self._backends]

    def embed(self, text: str) -> list[float]:
        if self._active is None:
            return [0.0] * self._target_dim
        vec = self._active.embed(text)
        if self._reducer:
            return self._reducer.reduce(vec)
        return _normalize_dimension(vec, self._target_dim)

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        if self._active is None:
            return [[0.0] * self._target_dim for _ in texts]
        vectors = self._active.embed_batch(texts)
        if self._reducer:
            return [self._reducer.reduce(v) for v in vectors]
        return [_normalize_dimension(v, self._target_dim) for v in vectors]

    def benchmark(self, texts: list[str]) -> dict[str, dict[str, float | str]]:
        """Benchmark all available backends."""
        results: dict[str, dict[str, float | str]] = {}
        sample_count = max(1, len(texts))
        for backend in self._backends:
            start = time.time()
            try:
                for text in texts:
                    backend.embed(text)
                elapsed = time.time() - start
                results[backend.name] = {
                    "total_ms": float(elapsed * 1000.0),
                    "per_text_ms": float((elapsed / sample_count) * 1000.0),
                    "dimension": float(backend.dimension),
                }
            except Exception as exc:
                results[backend.name] = {"error": str(exc)}
        return results

    def _initialize(self) -> None:
        preferred = str(self._config.get("preferred_backend", "auto")).strip().lower()
        order: list[str] = []
        if preferred and preferred != "auto":
            order.append(preferred)
        for name in ("sentence-transformers", "ollama", "openai-compat", "tfidf"):
            if name not in order:
                order.append(name)

        for name in order:
            backend = self._build_backend(name)
            if backend is not None:
                self._backends.append(backend)

        if not self._backends:
            self._backends.append(TFIDFBackend(dimension=self._target_dim))

        self._active = self._backends[0]
        if self._active.dimension != self._target_dim:
            self._reducer = PCAReducer(self._target_dim)

        logger.info(
            "Embedding router active backend=%s (dim=%d->%d)",
            self._active.name,
            self._active.dimension,
            self._target_dim,
        )

    def _build_backend(self, name: str) -> EmbeddingBackend | None:
        canonical = name.replace("_", "-")

        if canonical == "tfidf":
            return TFIDFBackend(dimension=self._target_dim)

        if canonical == "sentence-transformers":
            try:
                model = str(self._config.get("st_model", "all-MiniLM-L6-v2"))
                return SentenceTransformerBackend(model)
            except Exception:
                return None

        if canonical == "ollama":
            try:
                model = str(self._config.get("ollama_model", "nomic-embed-text"))
                base_url = str(self._config.get("ollama_base_url", "http://127.0.0.1:11434"))
                backend = OllamaEmbeddingBackend(model=model, base_url=base_url)
                is_available = getattr(getattr(backend, "_client", None), "is_available", None)
                if callable(is_available) and not is_available():
                    return None
                return backend
            except Exception:
                return None

        if canonical in {"openai-compat", "openai", "openai-compatible"}:
            try:
                return OpenAICompatBackend(
                    base_url=str(self._config.get("openai_base_url", "http://127.0.0.1:8080")),
                    api_key=str(self._config.get("openai_api_key", "")),
                    model=str(self._config.get("openai_model", "text-embedding-3-small")),
                    timeout_s=float(self._config.get("openai_timeout_s", 10.0)),
                )
            except Exception:
                return None

        return None
