"""Hybrid retriever: BM25 (in-memory index over course_chunks) + optional embeddings."""

from __future__ import annotations

import asyncio
import math
from collections import Counter, defaultdict
from dataclasses import dataclass

import numpy as np
from sqlalchemy import select

from app.core.logging import get_logger
from app.db import SessionLocal
from app.models import CourseChunk
from app.rag.chunking import tokenize

log = get_logger("rag.retriever")


@dataclass
class Hit:
    chunk_id: str
    course: str
    module: str
    section: str
    page: str
    url: str
    content: str
    score: float
    visibility: str

    def source_label(self) -> str:
        sec = f", {self.section}" if self.section and self.section != self.module else ""
        return f"{self.course} – {self.module}{sec}"


class BM25Index:
    def __init__(self, k1: float = 1.4, b: float = 0.75) -> None:
        self.k1, self.b = k1, b
        self.docs: list[dict] = []
        self.tf: list[Counter[str]] = []
        self.df: Counter[str] = Counter()
        self.doc_len: list[int] = []
        self.avg_len = 1.0
        self.embeddings: np.ndarray | None = None
        self.postings: dict[str, list[int]] = defaultdict(list)

    def build(self, rows: list[CourseChunk]) -> None:
        self.docs, self.tf, self.doc_len = [], [], []
        self.df = Counter()
        self.postings = defaultdict(list)
        embs: list[list[float]] = []
        have_emb = True
        for i, r in enumerate(rows):
            head = f"{r.module} {r.section} "
            toks = tokenize(head * 2 + r.content)
            tf = Counter(toks)
            self.docs.append({"id": r.id, "course": r.course_code, "module": r.module,
                              "section": r.section, "page": r.page, "url": r.url,
                              "content": r.content, "visibility": r.visibility})
            self.tf.append(tf)
            self.doc_len.append(len(toks))
            for t in tf:
                self.df[t] += 1
                self.postings[t].append(i)
            if r.embedding:
                embs.append(r.embedding)
            else:
                have_emb = False
        self.avg_len = (sum(self.doc_len) / len(self.doc_len)) if self.doc_len else 1.0
        self.embeddings = np.array(embs, dtype=np.float32) if (have_emb and embs) else None
        log.info("bm25_built", docs=len(self.docs), embeddings=self.embeddings is not None)

    def _idf(self, t: str) -> float:
        n = len(self.docs)
        df = self.df.get(t, 0)
        return math.log(1 + (n - df + 0.5) / (df + 0.5))

    def search(self, query: str, top_k: int = 8, courses: set[str] | None = None,
               boost_course: str | None = None, include_professor: bool = False,
               query_embedding: list[float] | None = None) -> list[Hit]:
        q = tokenize(query)
        if not q or not self.docs:
            return []
        scores: dict[int, float] = defaultdict(float)
        for t in set(q):
            idf = self._idf(t)
            for i in self.postings.get(t, ()):
                f = self.tf[i][t]
                denom = f + self.k1 * (1 - self.b + self.b * self.doc_len[i] / self.avg_len)
                scores[i] += idf * f * (self.k1 + 1) / denom
        if scores:
            mx = max(scores.values()) or 1.0
            for i in scores:
                scores[i] /= mx
        if query_embedding is not None and self.embeddings is not None:
            qv = np.array(query_embedding, dtype=np.float32)
            sims = self.embeddings @ qv / (np.linalg.norm(self.embeddings, axis=1) *
                                           (np.linalg.norm(qv) or 1.0) + 1e-9)
            top = np.argsort(-sims)[: top_k * 4]
            for i in top:
                scores[int(i)] = 0.35 * scores.get(int(i), 0.0) + 0.65 * float(sims[i])
        hits: list[Hit] = []
        for i, s in scores.items():
            d = self.docs[i]
            if courses and d["course"] not in courses:
                continue
            if d["visibility"] == "professor_only" and not include_professor:
                continue
            if boost_course and d["course"] == boost_course:
                s *= 1.15
            hits.append(Hit(d["id"], d["course"], d["module"], d["section"], d["page"],
                            d["url"], d["content"], s, d["visibility"]))
        hits.sort(key=lambda h: h.score, reverse=True)
        # diversify: at most 3 hits per module
        out: list[Hit] = []
        per_module: Counter[str] = Counter()
        for h in hits:
            if per_module[h.module] >= 3:
                continue
            per_module[h.module] += 1
            out.append(h)
            if len(out) >= top_k:
                break
        return out


index = BM25Index()
_lock = asyncio.Lock()


async def rebuild_index() -> int:
    async with _lock:
        async with SessionLocal() as session:
            rows = list((await session.execute(select(CourseChunk))).scalars())
        index.build(rows)
        return len(rows)


def format_for_model(hits: list[Hit]) -> str:
    if not hits:
        return "Aucun passage pertinent trouvé dans le matériel du cours."
    parts = ["Passages du matériel de cours (cite-les avec l'identifiant entre crochets) :"]
    for i, h in enumerate(hits, 1):
        parts.append(f"[S{i}] Source: {h.source_label()}"
                     f"{' — ' + h.url if h.url else ''}\n<document>\n{h.content[:3000]}\n</document>")
    return "\n\n".join(parts)