"""Hybrid retrieval: BM25 + dense vectors, fused with reciprocal rank fusion (lessons 06 and 08).

Dense vectors come from an Embedder. The default, LsaEmbedder (TF-IDF → truncated SVD), needs no model
download or API key. OpenAIEmbedder uses text-embedding-3-small and is what you would run in production.
"""

from __future__ import annotations

import math
import re
from collections import Counter
from dataclasses import dataclass
from typing import Protocol

import numpy as np

from .ingest import Chunk

TOKEN = re.compile(r"[a-z0-9_]+")
STOP = frozenset("a an the and or of to in on for is are was be by with as at it this that from how do i "
                 "my what when why does doesn t s can you your we".split())


def tokenize(text: str) -> list[str]:
    return [t for t in TOKEN.findall(text.lower()) if t not in STOP]


class BM25:
    def __init__(self, texts: list[str], k1: float = 1.5, b: float = 0.75) -> None:
        self.k1, self.b = k1, b
        docs = [Counter(tokenize(t)) for t in texts]
        self.lengths = np.array([sum(d.values()) for d in docs], dtype=float)
        self.avg_len = float(self.lengths.mean())
        df = Counter(term for d in docs for term in d)
        n = len(docs)
        self.idf = {t: math.log(1 + (n - f + 0.5) / (f + 0.5)) for t, f in df.items()}
        self.postings: dict[str, list[tuple[int, int]]] = {}
        for i, d in enumerate(docs):
            for term, tf in d.items():
                self.postings.setdefault(term, []).append((i, tf))

    def scores(self, query: str) -> np.ndarray:
        s = np.zeros(len(self.lengths))
        for term in set(tokenize(query)):
            for i, tf in self.postings.get(term, ()):
                norm = 1 - self.b + self.b * self.lengths[i] / self.avg_len
                s[i] += self.idf[term] * tf * (self.k1 + 1) / (tf + self.k1 * norm)
        return s


class Embedder(Protocol):
    name: str

    def fit(self, texts: list[str]) -> np.ndarray: ...
    def embed(self, texts: list[str]) -> np.ndarray: ...


def _unit(m: np.ndarray) -> np.ndarray:
    return m / np.maximum(np.linalg.norm(m, axis=1, keepdims=True), 1e-12)


class LsaEmbedder:
    """Offline dense vectors: TF-IDF then SVD. A stand-in for a neural embedding model."""
    name = "lsa-256"

    def __init__(self, dims: int = 256, seed: int = 0) -> None:
        from sklearn.decomposition import TruncatedSVD
        from sklearn.feature_extraction.text import TfidfVectorizer
        self.tfidf = TfidfVectorizer(tokenizer=tokenize, lowercase=False, token_pattern=None,
                                     sublinear_tf=True, min_df=2)
        self.svd = TruncatedSVD(n_components=dims, random_state=seed)

    def fit(self, texts: list[str]) -> np.ndarray:
        return _unit(self.svd.fit_transform(self.tfidf.fit_transform(texts)))

    def embed(self, texts: list[str]) -> np.ndarray:
        return _unit(self.svd.transform(self.tfidf.transform(texts)))


class OpenAIEmbedder:
    """Real embeddings (needs OPENAI_API_KEY). Batches requests; vectors are already unit length."""
    name = "text-embedding-3-small"

    def __init__(self, client, batch_size: int = 100) -> None:
        self.client, self.batch_size = client, batch_size

    def embed(self, texts: list[str]) -> np.ndarray:
        vectors = []
        for i in range(0, len(texts), self.batch_size):
            resp = self.client.embeddings.create(model=self.name, input=texts[i:i + self.batch_size])
            vectors += [d.embedding for d in resp.data]
        return _unit(np.array(vectors, dtype=np.float32))

    fit = embed


def rrf(*rankings: list[int], k: int = 60) -> list[tuple[int, float]]:
    fused: Counter = Counter()
    for ranking in rankings:
        for rank, i in enumerate(ranking, start=1):
            fused[i] += 1 / (k + rank)
    return fused.most_common()


@dataclass(frozen=True)
class Hit:
    chunk: Chunk
    score: float
    rank: int


class HybridIndex:
    def __init__(self, chunks: list[Chunk], embedder: Embedder | None = None) -> None:
        self.chunks = chunks
        texts = [f"{c.header()}\n{c.text}" for c in chunks]
        self.embedder = embedder or LsaEmbedder()
        self.bm25 = BM25(texts)
        self.vectors = self.embedder.fit(texts)

    def __len__(self) -> int:
        return len(self.chunks)

    def search(self, query: str, k: int = 5, course: str | None = None, mode: str = "hybrid",
               depth: int = 60) -> list[Hit]:
        allowed = None if course is None else {i for i, c in enumerate(self.chunks) if c.course == course}
        bm = self._top(self.bm25.scores(query), depth, allowed)
        if mode == "bm25":
            ranked = [(i, 1 / (60 + r)) for r, i in enumerate(bm, start=1)]
        else:
            dense = self._top(self.vectors @ self.embedder.embed([query])[0], depth, allowed)
            ranked = ([(i, 1 / (60 + r)) for r, i in enumerate(dense, start=1)] if mode == "dense"
                      else rrf(bm, dense))
        return [Hit(self.chunks[i], round(s, 6), r) for r, (i, s) in enumerate(ranked[:k], start=1)]

    @staticmethod
    def _top(scores: np.ndarray, depth: int, allowed: set[int] | None) -> list[int]:
        if allowed is not None:
            mask = np.full(len(scores), -np.inf)
            mask[list(allowed)] = 0
            scores = scores + mask
        top = np.argsort(-scores)[:depth]
        return [int(i) for i in top if np.isfinite(scores[i])]