"""Tiny retriever under test: heading-aware chunks + BM25. Config is a dataclass so a 'regression' is one changed field."""
import re, glob, math, collections, os
from dataclasses import dataclass

CORPUS = os.path.join(os.path.dirname(__file__), "..", "corpus", "*.md")
TOK = re.compile(r"[a-z0-9_]+")
HEADING = re.compile(r"^(#{1,3}) (.+?)\s*(\{.*\})?$", re.M)

@dataclass(frozen=True)
class Config:
    chunk_size: int = 800        # target characters per chunk
    heading_prefix: bool = True  # index "H1 > H2 > H3" together with the body
    lowercase: bool = True

def toks(s, cfg):
    return TOK.findall(s.lower() if cfg.lowercase else s)

def sections(text):
    ms = list(HEADING.finditer(text)); out = []; stack = []
    for i, m in enumerate(ms):
        lvl = len(m.group(1)); title = m.group(2).strip()
        stack = [s for s in stack if s[0] < lvl] + [(lvl, title)]
        end = ms[i + 1].start() if i + 1 < len(ms) else len(text)
        out.append((m.start(), end, " > ".join(t for _, t in stack)))
    return out

def paragraphs(text, size):
    out, start, cur = [], 0, 0
    for m in re.finditer(r"\n\s*\n", text):
        if m.end() - start > size and cur > start:
            out.append((start, cur)); start = cur
        cur = m.end()
    out.append((start, len(text)))
    return [c for c in out if text[c[0]:c[1]].strip()]

def chunk(text, cfg):
    secs = sections(text) or [(0, len(text), "")]
    if secs[0][0] > 0: secs = [(0, secs[0][0], "")] + secs
    out = []
    for s, e, path in secs:
        for a, b in paragraphs(text[s:e], cfg.chunk_size):
            body = text[s + a:s + b]
            out.append((f"{path}\n{body}" if cfg.heading_prefix and path else body))
    return out

class Retriever:
    def __init__(self, cfg=Config()):
        self.cfg = cfg
        self.doc, self.tf = [], []
        for p in sorted(glob.glob(CORPUS)):
            name = os.path.basename(p)
            for c in chunk(open(p, encoding="utf-8").read(), cfg):
                self.doc.append(name); self.tf.append(collections.Counter(toks(c, cfg)))
        n = len(self.tf)
        df = collections.Counter(t for c in self.tf for t in c)
        self.L = [sum(c.values()) for c in self.tf]; self.avg = sum(self.L) / n
        self.idf = {t: math.log(1 + (n - k + .5) / (k + .5)) for t, k in df.items()}
        self.inv = collections.defaultdict(list)
        for i, c in enumerate(self.tf):
            for t, f in c.items(): self.inv[t].append((i, f))

    def search(self, q, k=5):
        """Return top-k *documents* (best chunk per doc), ranked."""
        sc = collections.defaultdict(float)
        for t in set(toks(q, self.cfg)):
            for i, f in self.inv.get(t, ()):
                sc[i] += self.idf[t] * f * 2.5 / (f + 1.5 * (.25 + .75 * self.L[i] / self.avg))
        seen, out = set(), []
        for i, _ in sorted(sc.items(), key=lambda x: -x[1]):
            if self.doc[i] not in seen:
                seen.add(self.doc[i]); out.append(self.doc[i])
                if len(out) == k: break
        return out