"""Embedding- und Rerank-Clients sowie Skill-Index."""

from __future__ import annotations

import hashlib
import math
import re
from collections.abc import Sequence
from typing import Protocol

import httpx

Vector = list[float]


class Embedder(Protocol):
    def embed(self, texts: Sequence[str]) -> list[Vector]: ...


class Reranker(Protocol):
    def rerank(self, query: str, documents: Sequence[str]) -> list[float]: ...


def cosine(a: Sequence[float], b: Sequence[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b, strict=True))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)


class HashingEmbedder:
    """Deterministischer Bag-of-Words-Embedder für Tests und als Notfall-Fallback ohne GPU/Netz."""

    _WORD = re.compile(r"\w+", re.UNICODE)

    def __init__(self, dim: int = 512):
        self.dim = dim

    def _vector(self, text: str) -> Vector:
        vec = [0.0] * self.dim
        for word in self._WORD.findall(text.lower()):
            h = int.from_bytes(hashlib.sha1(word.encode()).digest()[:4], "big")
            vec[h % self.dim] += 1.0 if (h >> 31) & 1 else -1.0
        norm = math.sqrt(sum(v * v for v in vec)) or 1.0
        return [v / norm for v in vec]

    def embed(self, texts: Sequence[str]) -> list[Vector]:
        return [self._vector(t) for t in texts]


class OpenAIEmbedder:
    """Embeddings über ein OpenAI-kompatibles API (LiteLLM → llama.cpp/TEI)."""

    def __init__(self, base_url: str, model: str, api_key: str = "", timeout: float = 30.0,
                 client: httpx.Client | None = None):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
        self.client = client or httpx.Client(timeout=timeout)

    def embed(self, texts: Sequence[str]) -> list[Vector]:
        resp = self.client.post(
            f"{self.base_url}/v1/embeddings",
            json={"model": self.model, "input": list(texts)},
            headers=self.headers,
        )
        resp.raise_for_status()
        data = sorted(resp.json()["data"], key=lambda d: d["index"])
        return [d["embedding"] for d in data]


class LlamaCppReranker:
    """Reranking über den llama.cpp-Server (`--reranking`, Endpoint /v1/rerank)."""

    def __init__(self, base_url: str, model: str = "rerank-default", timeout: float = 30.0,
                 client: httpx.Client | None = None):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.client = client or httpx.Client(timeout=timeout)

    def rerank(self, query: str, documents: Sequence[str]) -> list[float]:
        resp = self.client.post(
            f"{self.base_url}/v1/rerank",
            json={"model": self.model, "query": query, "documents": list(documents)},
        )
        resp.raise_for_status()
        scores = [0.0] * len(documents)
        for item in resp.json()["results"]:
            scores[item["index"]] = float(item["relevance_score"])
        return scores


class SkillIndex:
    """Hält pro Skill die Vektoren von Beschreibung und Triggern (wenige hundert Skills → im RAM)."""

    def __init__(self, embedder: Embedder):
        self.embedder = embedder
        self._cache: dict[str, Vector] = {}
        self._entries: dict[str, list[Vector]] = {}

    def _embed_cached(self, texts: list[str]) -> list[Vector]:
        missing = [t for t in texts if t not in self._cache]
        if missing:
            for text, vec in zip(missing, self.embedder.embed(missing), strict=True):
                self._cache[text] = vec
        return [self._cache[t] for t in texts]

    def rebuild(self, skills: dict[str, list[str]]) -> None:
        """skills: name -> Texte (Beschreibung + Trigger)."""
        all_texts = sorted({t for texts in skills.values() for t in texts})
        self._embed_cached(all_texts)
        self._entries = {name: [self._cache[t] for t in texts] for name, texts in skills.items()}

    def vector(self, name: str) -> Vector:
        return self._entries[name][0]

    def search(self, query_vec: Vector, k: int = 5, only: set[str] | None = None) -> list[tuple[str, float]]:
        scored = []
        for name, vecs in self._entries.items():
            if only is not None and name not in only:
                continue
            scored.append((name, max(cosine(query_vec, v) for v in vecs)))
        scored.sort(key=lambda x: (-x[1], x[0]))
        return scored[:k]

    def pairwise_similar(self, threshold: float) -> list[tuple[str, str, float]]:
        names = sorted(self._entries)
        pairs = []
        for i, a in enumerate(names):
            for b in names[i + 1:]:
                sim = cosine(self.vector(a), self.vector(b))
                if sim >= threshold:
                    pairs.append((a, b, sim))
        return pairs