#!/usr/bin/env python3
"""Ordonarea chunk-urilor: embeddings + potrivire lexicala, fuzionate.

De ce nu doar embeddings. Masurat pe indexul viu (140 de chunk-uri), scorul
cosinus al celui mai bun chunk a fost 0,600-0,816 la intrebari care CHIAR au
raspuns in documente si 0,534-0,685 la intrebari straine de ROA. Intervalele SE
SUPRAPUN, deci niciun prag pe cosinus singur nu le separa: „bilant contabil"
(0,600, buna) sta sub „eroare 0x80070005 la instalarea unui joc pe Steam" (0,685,
straina).

Ce lipseste din cosinus e exact ce conteaza la suport: **codurile**. `ORA-01722`,
`D406`, `CIF`, `SPV` sunt tokeni exacti — o cautare lexicala ii prinde imediat, iar
una semantica ii topeste in "ceva despre facturi". Asa ca ordonam de doua ori,
independent, si fuzionam clasamentele (Reciprocal Rank Fusion): un chunk urcat de
ambele metode iese primul, iar unul urcat doar de una ramane in cursa.

Fuziunea da mereu un clasament, si la o intrebare complet straina — de aceea
decizia "avem sau nu acoperire in documente" se ia separat, in `assess()`, pe
dovezi (cat de sus e cosinusul, si daca termenii rari din intrebare chiar apar in
chunk), nu pe pozitia in clasament.
"""

from __future__ import annotations

import math
import re
import unicodedata

import config

# Praguri implicite, calibrate pe masuratoarea din docstring. Se pot suprascrie
# din env (vezi config.DEFAULTS) fara sa umbli in cod.
RRF_K = 60                # constanta standard din Reciprocal Rank Fusion


def strong_cosine() -> float:
    """Peste atat, singur cosinusul e de ajuns (RANK_STRONG_COSINE)."""
    return float(config.get("RANK_STRONG_COSINE", "0.70"))


def weak_cosine() -> float:
    """Sub atat nu salvam raspunsul nici cu potrivire lexicala (RANK_WEAK_COSINE)."""
    return float(config.get("RANK_WEAK_COSINE", "0.58"))


def recall_n() -> int:
    """Cate chunk-uri intra in fuziune, din fiecare metoda (RANK_RECALL_N)."""
    return config.get_int("RANK_RECALL_N", 12)


def min_rare_ratio() -> float:
    """Ce fractiune din termenii distinctivi ai intrebarii trebuie gasita (RANK_MIN_RARE_RATIO).

    Un singur termen gasit nu e o dovada: „cum imi resetez parola de la Windows" a
    fost salvata gresit doar fiindca „parola" apare in documente. Cerand o
    fractiune, intrebarea trebuie sa se suprapuna cu documentele, nu sa le atinga.
    """
    return float(config.get("RANK_MIN_RARE_RATIO", "0.5"))

_TOKEN = re.compile(r"[a-z0-9]+(?:-[a-z0-9]+)*")

# Cuvinte prea frecvente ca sa spuna ceva despre subiect. Lista e scurta
# deliberat: la BM25 idf-ul se ocupa oricum de restul.
_STOP = {
    "si", "sa", "se", "la", "de", "din", "in", "pe", "cu", "pentru", "ca", "care",
    "este", "sunt", "am", "are", "nu", "un", "o", "al", "ale", "lui", "cum", "ce",
    "mi", "imi", "ma", "eu", "va", "vrea", "vreau", "poate", "pot", "the", "to",
    "a", "an", "of", "is", "it", "da", "dar", "sau", "mai", "fost", "face",
}

# Coduri de eroare si formulare: valoreaza cat zece cuvinte obisnuite.
_COD = re.compile(r"^(?:ora|pls|tns|sp2|imp|exp|ora)-\d+$|^d\d{3}$|^e\d{3,}$")


def fold(text: str) -> str:
    """Fara diacritice, litere mici. „Factură" si „factura" trebuie sa fie acelasi token."""
    norm = unicodedata.normalize("NFKD", text.lower())
    return "".join(c for c in norm if not unicodedata.combining(c))


def tokenize(text: str) -> list[str]:
    return [t for t in _TOKEN.findall(fold(text)) if t not in _STOP and len(t) > 1]


class Bm25:
    """Okapi BM25 peste chunk-urile indexului.

    Se reconstruieste la fiecare recitire a indexului (la 140 de chunk-uri costa
    milisecunde) — nu merita tinut in fisier si desincronizat de el.
    """

    def __init__(self, documents: list[str], k1: float = 1.5, b: float = 0.75):
        self.k1, self.b = k1, b
        self.docs = [tokenize(d) for d in documents]
        self.lens = [len(d) for d in self.docs]
        self.avg_len = (sum(self.lens) / len(self.lens)) if self.lens else 0.0
        self.freqs: list[dict[str, int]] = []
        self.df: dict[str, int] = {}
        for toks in self.docs:
            f: dict[str, int] = {}
            for t in toks:
                f[t] = f.get(t, 0) + 1
            self.freqs.append(f)
            for t in f:
                self.df[t] = self.df.get(t, 0) + 1
        self.n = len(self.docs)

    def idf(self, term: str) -> float:
        df = self.df.get(term, 0)
        # varianta cu +0.5 (Robertson): termenii din peste jumatate din documente
        # ajung aproape de zero, ceea ce e exact ce vrem.
        return math.log(1 + (self.n - df + 0.5) / (df + 0.5))

    def scores(self, query: str) -> list[float]:
        q = tokenize(query)
        out = [0.0] * self.n
        if not q or not self.avg_len:
            return out
        for i, f in enumerate(self.freqs):
            s = 0.0
            for t in q:
                tf = f.get(t, 0)
                if not tf:
                    continue
                denom = tf + self.k1 * (1 - self.b + self.b * self.lens[i] / self.avg_len)
                s += self.idf(t) * tf * (self.k1 + 1) / denom
            out[i] = s
        return out

    def rare_terms(self, query: str, max_df_ratio: float = 0.25) -> list[str]:
        """Termenii din intrebare care chiar disting ceva: coduri, sau cuvinte rare in corpus."""
        out = []
        for t in dict.fromkeys(tokenize(query)):
            if _COD.match(t):
                out.append(t)
            elif len(t) >= 4 and self.df.get(t, 0) <= max(1, int(self.n * max_df_ratio)):
                out.append(t)
        return out


def rrf(rankings: list[list[int]], k: int = RRF_K) -> dict[int, float]:
    """Reciprocal Rank Fusion: suma de 1/(k+pozitie) peste clasamente.

    Fuzioneaza POZITII, nu scoruri — cosinusul (0..1) si BM25 (nemarginit) nu se
    pot aduna direct fara o normalizare care ar depinde de corpus.
    """
    fused: dict[int, float] = {}
    for ranking in rankings:
        for pos, idx in enumerate(ranking):
            fused[idx] = fused.get(idx, 0.0) + 1.0 / (k + pos + 1)
    return fused


def rank(
    cosines: list[float],
    bm25_scores: list[float],
    recall: int | None = None,
) -> list[int]:
    """Indicii chunk-urilor, de la cel mai relevant, dupa fuziunea celor doua clasamente."""
    recall = recall_n() if recall is None else recall
    by_cos = sorted(range(len(cosines)), key=lambda i: cosines[i], reverse=True)[:recall]
    # chunk-urile cu scor BM25 zero nu au niciun token comun cu intrebarea: nu au
    # ce cauta in clasamentul lexical, doar ar dilua fuziunea cu pozitii arbitrare.
    by_bm = [i for i in sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)
             if bm25_scores[i] > 0][:recall]
    fused = rrf([by_cos, by_bm])
    return sorted(fused, key=lambda i: (fused[i], cosines[i]), reverse=True)


def assess(
    query: str,
    best_cosine: float,
    top_texts: list[str],
    bm25: Bm25,
    *,
    strong: float | None = None,
    weak: float | None = None,
    min_ratio: float | None = None,
) -> tuple[bool, str]:
    """(avem acoperire in documente?, motivul) — decizia de a raspunde sau de a escalada.

    Trei cazuri, in ordine:
      1. cosinus mare  -> raspundem, indiferent de lexical;
      2. cosinus la mijloc, dar o fractiune suficienta din termenii rari ai
         intrebarii apare chiar in chunk-uri (tipic: un cod de eroare) -> raspundem;
      3. altfel -> nu avem acoperire, intrebarea pleaca la suport.
    """
    strong = strong_cosine() if strong is None else strong
    weak = weak_cosine() if weak is None else weak
    min_ratio = min_rare_ratio() if min_ratio is None else min_ratio
    if best_cosine >= strong:
        return True, f"cosinus {best_cosine:.3f} >= {strong}"
    if best_cosine < weak:
        return False, f"cosinus {best_cosine:.3f} < {weak}"
    rare = bm25.rare_terms(query)
    if not rare:
        return False, f"cosinus {best_cosine:.3f} la limita, iar intrebarea nu are termeni distinctivi"
    blob = fold("\n".join(top_texts))
    gasiti = [t for t in rare if t in blob]
    if gasiti and len(gasiti) / len(rare) >= min_ratio:
        return True, (
            f"cosinus {best_cosine:.3f} + {len(gasiti)}/{len(rare)} termeni distinctivi "
            f"in context: {', '.join(gasiti)}"
        )
    lipsa = [t for t in rare if t not in blob]
    return False, (
        f"cosinus {best_cosine:.3f} la limita, iar termenii distinctivi "
        f"({', '.join(lipsa[:5])}) nu apar in documente"
    )