"""Query-Conditioned Compressive Retrieval (QCCR).

Motivation
----------
Two failure modes of fragment-level selection at tight budgets:

  1. *Wrong chunk of the right file.* When a file is indexed as N fragments
     of ~400 tokens each, selecting fragment #3 may miss the paragraph in
     fragment #7 that actually answers the query.
  2. *Diversity spread.* Submodular log-det or MMR at the fragment level
     rewards picking many hash-distinct tiny chunks over one focused file.
     At 4K budget the answer drowns in uniform-weight filler.

QCCR sidesteps both by operating at two granularities:

  - *File-level BM25* chooses which documents are candidates. This is
    coarse, fast, and mirrors what a human does ("which file has this?").
  - *Sentence-level query-conditioned BM25 + MMR* (Carbonell-Goldstein 1998)
    picks the specific sentences from those files that answer the query,
    with diversity control against redundant sentences.

The result is a per-query extractive summary: ~15 tokens per sentence
× ~270 sentences = full 4K budget spent on content that directly
addresses the query, drawn from whichever files BM25 ranked highest.

Nothing here uses embeddings, neural inference, or trained weights.
Pure classical IR (BM25 + MMR) applied at the right granularity.

References
----------
- Robertson & Zaragoza (2009), "The Probabilistic Relevance Framework: BM25
  and Beyond"
- Carbonell & Goldstein (1998), "The Use of MMR, Diversity-Based Reranking
  for Reordering Documents and Producing Summaries"
- Nemhauser, Wolsey, Fisher (1978), (1-1/e) bound on greedy submodular
  maximization — applies to MMR with monotone-submodular relevance.
"""
from __future__ import annotations

import math
import re
from collections import Counter

# ── Tokenization (shared with dopt_selector style) ──────────────────────
_IDENT_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]{2,}")
_CAMEL_RE = re.compile(r"[A-Z]?[a-z]+|[A-Z]+(?=[A-Z]|$)|\d+")
_SENTENCE_SPLIT = re.compile(r"(?<=[.!?])\s+|(?<=\n\n)|(?<=;\n)|(?<=\}\n)")

_STOPWORDS = frozenset("""
the a an of to in on for with by how does what why is are and or but
do between two include shown actual actually that this these those
from can could should would has have had will was were been being
about into through over under above below after before while when where
not no yes all any some most many much more less than then thus therefore
you your their his her its our it he she they them we us i me my
""".split())

_BM25_K1 = 1.5
_BM25_B = 0.75
_MMR_LAMBDA = 0.7           # relevance-diversity tradeoff (70% relevance)
_MIN_SENTENCE_CHARS = 20    # skip shorter "sentences" (whitespace artifacts)
_MAX_FILES_CONSIDERED = 12  # top-K files per query
_CHARS_PER_TOKEN = 4        # approximation for token-budget accounting


def _split_identifier(tok: str) -> list[str]:
    """`taint_flow_total` → [taint_flow_total, taint, flow, total]. Also handles
    CamelCase. Improves recall when a query uses words that code identifiers
    concatenate."""
    low = tok.lower()
    parts = {low}
    for piece in low.split("_"):
        if len(piece) > 2:
            parts.add(piece)
    for piece in _CAMEL_RE.findall(tok):
        p = piece.lower()
        if len(p) > 2:
            parts.add(p)
    return [p for p in parts if p not in _STOPWORDS]


def _tokenize(text: str) -> list[str]:
    out: list[str] = []
    for raw in _IDENT_RE.findall(text):
        out.extend(_split_identifier(raw))
    return out


def _query_tokens(query: str) -> frozenset[str]:
    return frozenset(t for t in _tokenize(query) if len(t) > 2)


def _split_sentences(text: str) -> list[str]:
    """Split on sentence boundaries plus structural code breaks (blank line,
    semicolon+newline, close-brace+newline). Drops tiny fragments."""
    out: list[str] = []
    for chunk in _SENTENCE_SPLIT.split(text):
        s = chunk.strip()
        if len(s) >= _MIN_SENTENCE_CHARS:
            out.append(s)
    return out


def _approx_tokens(s: str) -> int:
    return max(1, len(s) // _CHARS_PER_TOKEN)


# ── BM25 scoring ─────────────────────────────────────────────────────────
def _bm25_corpus(texts: list[str]) -> tuple[list[Counter], list[int], Counter, float]:
    tf_list: list[Counter] = []
    lens: list[int] = []
    df: Counter = Counter()
    for t in texts:
        toks = _tokenize(t)
        tf = Counter(toks)
        tf_list.append(tf)
        lens.append(sum(tf.values()))
        for term in tf:
            df[term] += 1
    avgdl = (sum(lens) / len(lens)) if lens else 1.0
    return tf_list, lens, df, max(avgdl, 1.0)


def _bm25_score(
    q_terms: frozenset[str],
    tf: Counter,
    dl: int,
    df: Counter,
    N: int,
    avgdl: float,
) -> float:
    score = 0.0
    for term in q_terms:
        f = tf.get(term, 0)
        if f == 0:
            continue
        n = df.get(term, 0)
        idf = math.log(1.0 + (N - n + 0.5) / (n + 0.5))
        norm = 1.0 - _BM25_B + _BM25_B * (dl / avgdl)
        score += idf * (f * (_BM25_K1 + 1.0)) / (f + _BM25_K1 * norm)
    return score


# ── MMR sentence selection ──────────────────────────────────────────────
def _mmr_select(
    sentences: list[str],
    tf_list: list[Counter],
    rel: list[float],
    budget_tokens: int,
    lam: float = _MMR_LAMBDA,
) -> list[int]:
    """Maximum Marginal Relevance: balance query relevance against redundancy
    with already-selected sentences. Jaccard over token sets is the
    similarity surrogate — cheap and discrete, which is fine since sentence
    token sets are small. Returns selected indices in original document
    order.

        MMR = argmax_i [ λ · rel(i) − (1-λ) · max_{j∈S} sim(i, j) ]

    (Carbonell-Goldstein 1998). Greedy; near-optimal under the usual
    submodularity assumptions.
    """
    n = len(sentences)
    if n == 0:
        return []

    selected: list[int] = []
    remaining: list[int] = [i for i in range(n) if rel[i] > 0.0]
    if not remaining:
        return []
    # token sets per sentence for Jaccard
    sets = [frozenset(tf_list[i].keys()) for i in range(n)]
    budget_used = 0

    while remaining and budget_used < budget_tokens:
        if not selected:
            best = max(remaining, key=lambda i: rel[i])
        else:
            def mmr_score(i: int) -> float:
                sim = 0.0
                for j in selected:
                    a, b = sets[i], sets[j]
                    if not a or not b:
                        continue
                    inter = len(a & b)
                    union = len(a | b)
                    if union == 0:
                        continue
                    sim = max(sim, inter / union)
                return lam * rel[i] - (1.0 - lam) * sim
            best = max(remaining, key=mmr_score)

        if rel[best] <= 0.0:
            break
        cost = _approx_tokens(sentences[best])
        if budget_used + cost > budget_tokens:
            remaining.remove(best)
            continue
        selected.append(best)
        remaining.remove(best)
        budget_used += cost

    return sorted(selected)


# ── Public entry ─────────────────────────────────────────────────────────
def select(
    fragments: list[dict],
    token_budget: int,
    query: str = "",
) -> list[dict]:
    """Query-Conditioned Compressive Retrieval.

    Pipeline:
      1. Group fragments by source file.
      2. File-level BM25 → rank files by query relevance.
      3. For each top-K file, sentence-level BM25 + MMR extracts the
         sentences that answer the query.
      4. Emit as pseudo-fragments: one dict per file containing the
         extracted sentences concatenated, preserving original order.

    Empty query ⇒ fall back to original fragment list (no compression).
    """
    if not fragments:
        return []
    if not query:
        return fragments

    q_terms = _query_tokens(query)
    if not q_terms:
        return fragments

    # Group by file
    by_file: dict[str, list[dict]] = {}
    for raw in fragments:
        src = raw.get("source", "") or ""
        by_file.setdefault(src, []).append(raw)

    # File-level BM25
    file_sources = list(by_file.keys())
    file_texts = ["\n".join((r.get("content") or "") for r in by_file[s]) for s in file_sources]
    tf_list, lens, df, avgdl = _bm25_corpus(file_texts)
    N = len(file_sources)
    file_scores: list[tuple[float, str, str]] = [
        (_bm25_score(q_terms, tf_list[i], lens[i], df, N, avgdl), file_sources[i], file_texts[i])
        for i in range(N)
    ]
    file_scores.sort(reverse=True)

    # engine_s6 edit-target rerank — delegated to the centralized
    # service so this surface stays a thin caller and the recall-safe
    # try/except lives in exactly one place. Window-local permutation
    # of the existing BM25 candidate set; tail preserved (recall floor);
    # budget allocation below still uses each file's original BM25
    # score, so this only changes the ORDER excerpts are emitted in.
    # Validation: see entroly/file_localizer.py module docstring.
    if N > 1 and any(s > 0 for s, _, _ in file_scores):
        from .file_localizer import localize_files
        files_map = dict(zip(file_sources, file_texts))
        bm25_order = [src for _, src, _ in file_scores]
        reranked = localize_files(
            files_map, query, k=len(bm25_order), base_ranked=bm25_order,
        )
        score_text_by_src = {src: (sc, txt)
                             for sc, src, txt in file_scores}
        file_scores = [(score_text_by_src[s][0], s,
                        score_text_by_src[s][1])
                       for s in reranked if s in score_text_by_src]

    top_files = [fs for fs in file_scores[:_MAX_FILES_CONSIDERED] if fs[0] > 0]
    if not top_files:
        return []

    # Split budget roughly in proportion to file BM25, with floor per file.
    total_score = sum(s for s, _, _ in top_files) or 1.0
    budget_left = int(token_budget)
    per_file_budget: dict[str, int] = {}
    for score, src, _ in top_files:
        share = int(token_budget * (score / total_score))
        per_file_budget[src] = max(share, 256)  # floor so small-scoring files still contribute a snippet

    # Sentence-level MMR per file
    output: list[dict] = []
    for score, src, text in top_files:
        if budget_left <= 0:
            break
        sentences = _split_sentences(text)
        if not sentences:
            continue
        s_tf, s_lens, s_df, s_avg = _bm25_corpus(sentences)
        s_N = len(sentences)
        rel = [
            _bm25_score(q_terms, s_tf[i], s_lens[i], s_df, s_N, s_avg)
            for i in range(s_N)
        ]
        file_budget = min(per_file_budget.get(src, 256), budget_left)
        chosen = _mmr_select(sentences, s_tf, rel, file_budget)
        if not chosen:
            continue
        excerpt = "\n".join(sentences[i] for i in chosen)
        tokens_used = _approx_tokens(excerpt)
        # Emit as a synthetic fragment preserving source attribution.
        output.append({
            "fragment_id": f"qccr::{src}",
            "source": src,
            "content": excerpt,
            "token_count": tokens_used,
        })
        budget_left -= tokens_used

    return output
