"""BRAG - Bayesian RAG.

Closed-form Bayesian orchestration for retrieval-augmented QA. The orchestrator itself makes no
LLM calls: every decision is a few arithmetic operations on conjugate (Beta) posteriors.

Actions (each is an LLM call made by the *executor*, never by the policy):
  R1 : read top-k retrieved passages -> answer or NONE, plus a follow-up search query
  R2 : retrieve with the follow-up query, read old + new passages -> answer or NONE   (needs R1)
  CB : m parallel closed-book samples -> ONE vote (their mode), reliability by agreement level

Relaxations
  R1' self-samples of one LLM are one source (vote = mode, reliability kappa_g by agreement g)
  R2' concentrated distractors: a wrong vote lands on a given rival w.p. eta (K_eff = 1/eta)
  R3' reliability hyperpriors are population-level, calibrated offline on labelled data, frozen
  R4' separate reliabilities per action type
  R5' control by information rate: a* = argmax_a D_a / c_a,  stop when  P_err * D_a* / c_a* <= lambda
Posterior (per-query reliabilities integrated out in closed form):
  log w_i = sum_t [lnB(a_t+n_ti, b_t+N_t-n_ti) - lnB(a_t, b_t+N_t)] - n_i ln eta ;  log w_residual = -ln eta
"""
from __future__ import annotations
import json, math, re, string, threading, time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Callable

import numpy as np
from scipy.special import betaln
from rank_bm25 import BM25Okapi

# ----------------------------------------------------------------- text utilities
_PUNCT = set(string.punctuation)
_STOP = set('a an the of in on at to for by with from and or is was were are be been who what which when '
            'where whose how did does do that this as its it his her their'.split())

def norm(s: str) -> str:
    s = (s or '').lower()
    s = ''.join(ch for ch in s if ch not in _PUNCT)
    s = re.sub(r'\b(a|an|the)\b', ' ', s)
    return ' '.join(s.split())

def is_none(a) -> bool:
    return a is None or norm(str(a)) in ('none', '', 'unknown', 'not stated', 'not mentioned')

def tokenize(s: str) -> list[str]:
    return [w for w in re.findall(r'\w+', s.lower()) if w not in _STOP]

def est_tokens(s: str) -> int:
    return max(1, len(s) // 4)

class Clusters:
    """Answer strings -> clusters (equal after normalisation, or token containment)."""
    def __init__(self):
        self.reps: list[str] = []
    def id(self, a: str) -> int:
        n = norm(a); nt = set(n.split())
        for i, r in enumerate(self.reps):
            rt = set(r.split())
            if n == r or (nt and rt and (nt <= rt or rt <= nt)):
                return i
        self.reps.append(n)
        return len(self.reps) - 1

# ----------------------------------------------------------------- retrieval
class Corpus:
    def __init__(self, passages: list[str]):
        if not passages:
            raise ValueError('corpus is empty')
        self.passages = passages
        self.bm25 = BM25Okapi([tokenize(p) or ['_'] for p in passages])

    @classmethod
    def from_json(cls, path: str) -> 'Corpus':
        rows = json.load(open(path))
        return cls([f"{r['title']}: {r['text']}" for r in rows])

    def search(self, query: str, k: int, exclude: set[int] = frozenset()) -> list[int]:
        sc = self.bm25.get_scores(tokenize(query) or ['_'])
        order = np.argsort(-sc)
        return [int(j) for j in order[:k + len(exclude)] if int(j) not in exclude][:k]

# ----------------------------------------------------------------- LLM executor
LLM = Callable[[list[dict], float, int], tuple[str, int, int]]   # messages, temperature, max_tokens

SYS = 'You are a precise question-answering component. Follow the output format exactly.'

def p_cb(q):
    return [{'role': 'system', 'content': SYS},
            {'role': 'user', 'content': f'Answer from your own knowledge. Give only a short answer (1-6 words), '
                                         f'always a best guess.\nQuestion: {q}\nAnswer:'}]

def _ctx(ps):
    return '\n'.join(f'[{i + 1}] {p}' for i, p in enumerate(ps))

def p_r1(q, ps):
    return [{'role': 'system', 'content': SYS},
            {'role': 'user', 'content':
             'Use ONLY the passages. Return JSON {"answer": "...", "followup": "..."} where answer is a short '
             'answer (1-6 words) if the passages TOGETHER answer the question, else "NONE"; followup is the single '
             'most useful next search query for the missing information (name the bridge entity if known). '
             f'Always give a followup.\n\nPassages:\n{_ctx(ps)}\n\nQuestion: {q}\nJSON:'}]

def p_r2(q, ps):
    return [{'role': 'system', 'content': SYS},
            {'role': 'user', 'content':
             'Use ONLY the passages; combine facts across passages for multi-hop questions. Give only a short '
             'answer (1-6 words), or NONE if the passages do not contain the answer.'
             f'\n\nPassages:\n{_ctx(ps)}\n\nQuestion: {q}\nAnswer:'}]

def clean_answer(t: str) -> str:
    t = (t or '').strip().strip('"').strip()
    t = re.sub(r'^(final\s+)?answer\s*:\s*', '', t, flags=re.I)
    return t.split('\n')[0].strip().strip('."')

def parse_r1(t: str) -> tuple[str, str]:
    m = re.search(r'\{.*\}', t or '', flags=re.S)
    if m:
        try:
            d = json.loads(m.group(0))
            return clean_answer(str(d.get('answer', 'NONE'))), str(d.get('followup', '')).strip()
        except json.JSONDecodeError:
            pass
    a = re.search(r'"?answer"?\s*[:=]\s*"([^"]*)"', t or '', flags=re.I)
    f = re.search(r'"?followup"?\s*[:=]\s*"([^"]*)"', t or '', flags=re.I)
    return (clean_answer(a.group(1)) if a else 'NONE'), (f.group(1).strip() if f else '')

def hf_llm(model: str, token: str | None, provider: str = 'auto', timeout: float = 60.0) -> LLM:
    """model: a Hugging Face model id (served by HF Inference Providers) or the URL of an
    OpenAI-compatible endpoint (e.g. a dedicated HF Inference Endpoint, TGI or vLLM server)."""
    from huggingface_hub import InferenceClient
    if model.startswith(('http://', 'https://')):
        client = InferenceClient(base_url=model, api_key=token, timeout=timeout)
    else:
        client = InferenceClient(model=model, token=token, provider=provider, timeout=timeout)
    def call(messages, temperature, max_tokens):
        r = client.chat_completion(messages, max_tokens=max_tokens, temperature=max(temperature, 0.01))
        text = r.choices[0].message.content or ''
        u = getattr(r, 'usage', None)
        tin = getattr(u, 'prompt_tokens', None) or sum(est_tokens(m['content']) for m in messages)
        tout = getattr(u, 'completion_tokens', None) or est_tokens(text)
        return text, int(tin), int(tout)
    return call

# ----------------------------------------------------------------- calibration
DEFAULT_CAL = {  # overwritten by data/calibration.json when present
    'rel': {'R1': [29, 6], 'R2': [54, 11]}, 'ans': {'R1': [34, 58], 'R2': [64, 28]},
    'kap': {'1': [2, 30], '2': [15, 24], '3': [14, 11]}, 'pg': {'1': 0.333, '2': 0.409, '3': 0.258},
    'eta': 0.591}

def _kl(rho, eta, q=1.0):
    rho = min(max(rho, 1e-6), 1 - 1e-6)
    return q * (rho - (1 - rho) * eta) * math.log(rho / ((1 - rho) * eta))

def gold_match(pred, gold: list[str]) -> int:
    if is_none(pred):
        return 0
    p = norm(pred)
    return int(any(norm(g) and (norm(g) in p or (len(p) > 3 and p in norm(g))) for g in gold))

# ----------------------------------------------------------------- the orchestrator
@dataclass
class Result:
    answer: str | None
    trace: list = field(default_factory=list)
    llm_calls: int = 0
    tokens_in: int = 0
    tokens_out: int = 0
    rounds: int = 0
    llm_seconds: float = 0.0
    policy_microseconds: float = 0.0
    stopped_because: str = ''
    def to_dict(self):
        return dict(answer=self.answer, stopped_because=self.stopped_because, llm_calls=self.llm_calls,
                    tokens_in=self.tokens_in, tokens_out=self.tokens_out, rounds=self.rounds,
                    llm_wall_seconds=round(self.llm_seconds, 3),
                    policy_microseconds_total=round(self.policy_microseconds, 1), trace=self.trace)

class BRAG:
    def __init__(self, llm: LLM, calibration: dict | None = None, lam: float = 0.05,
                 k: int = 5, cb_samples: int = 3, cb_temperature: float = 0.7):
        c = json.loads(json.dumps(calibration or DEFAULT_CAL))
        c['kap'] = {int(g): v for g, v in c['kap'].items()}; c['pg'] = {int(g): v for g, v in c['pg'].items()}
        self.P, self.llm, self.lam, self.k = c, llm, lam, k
        self.m, self.temp = cb_samples, cb_temperature

    # closed-form posterior over observed answer clusters + residual
    def posterior(self, votes):
        eta = self.P['eta']; cl = sorted({i for _, i in votes}); lw = {i: 0.0 for i in cl}
        by = {}
        for t, i in votes:
            by.setdefault(t, Counter())[i] += 1
        for t, cnt in by.items():
            a, b = self.P['kap'][t] if isinstance(t, int) else self.P['rel'][t]; N = sum(cnt.values())
            for i in cl:
                n = cnt.get(i, 0)
                lw[i] += betaln(a + n, b + N - n) - betaln(a, b + N) - n * math.log(eta)
        v = np.array([lw[i] for i in cl] + [-math.log(eta)]); v -= v.max(); p = np.exp(v); p /= p.sum()
        return cl, p

    def rates(self, done):
        eta = self.P['eta']; out = {}
        for t in ('R1', 'R2'):
            if t in done or (t == 'R2' and 'R1' not in done):
                continue
            a, b = self.P['ans'][t]; ra, rb = self.P['rel'][t]
            out[t] = _kl(ra / (ra + rb), eta, a / (a + b)) / 1.0
        if 'CB' not in done:
            d = sum(self.P['pg'][g] * _kl(self.P['kap'][g][0] / sum(self.P['kap'][g]), eta) for g in (1, 2, 3))
            out['CB'] = d / self.m
        return out

    _lock = threading.Lock()

    def _call(self, res, messages, temperature, max_tokens):
        t0 = time.perf_counter(); text, ti, to = self.llm(messages, temperature, max_tokens)
        with self._lock:
            res.llm_calls += 1
            res.tokens_in += ti; res.tokens_out += to
        return text

    def run(self, question: str, corpus: Corpus) -> Result:
        res = Result(answer=None); C = Clusters(); votes = []; raw = {}; done = set(); st = {}
        while True:
            t0 = time.perf_counter()
            cl, p = self.posterior(votes)
            pbest = float(p[:-1].max()) if cl else 0.0; perr = 1 - pbest
            rates = self.rates(done); act = None; why = ''
            if not rates:
                why = 'no actions left'
            else:
                a_star = max(rates, key=rates.get)
                if perr * rates[a_star] > self.lam:
                    act = a_star
                else:
                    why = f'P_err*D/c = {perr * rates[a_star]:.4f} <= lambda = {self.lam}'
            dt = 1e6 * (time.perf_counter() - t0); res.policy_microseconds += dt
            step = dict(posterior={(C.reps[i] if i < len(C.reps) else str(i)): round(float(pp), 4)
                                   for i, pp in zip(cl, p[:-1])},
                        residual=round(float(p[-1]), 4), p_error=round(perr, 4),
                        info_rate_per_call={k_: round(v, 4) for k_, v in rates.items()},
                        action=act or 'STOP', policy_us=round(dt, 1))
            if act is None:
                res.stopped_because = why; res.trace.append(step); break
            done.add(act); res.rounds += 1; tw = time.perf_counter()
            if act == 'CB':
                with ThreadPoolExecutor(self.m) as ex:
                    outs = list(ex.map(lambda _: self._call(res, p_cb(question), self.temp, 24), range(self.m)))
                samples = [clean_answer(o) for o in outs]; good = [s for s in samples if not is_none(s)]
                step['observation'] = {'samples': samples}
                if good:
                    ids = [C.id(s) for s in good]; mo, g = Counter(ids).most_common(1)[0]
                    votes.append((g, mo)); raw.setdefault(mo, good[ids.index(mo)])
                    step['observation']['vote'] = good[ids.index(mo)]; step['observation']['agreement'] = g
            elif act == 'R1':
                idx = corpus.search(question, self.k); st['idx1'] = idx
                ps = [corpus.passages[j][:600] for j in idx]
                ans, fu = parse_r1(self._call(res, p_r1(question, ps), 0.0, 96)); st['followup'] = fu or question
                step['observation'] = {'answer': ans, 'followup': fu, 'passages': [x[:80] for x in ps]}
                if not is_none(ans):
                    cid = C.id(ans); votes.append(('R1', cid)); raw.setdefault(cid, ans)
            else:  # R2
                idx2 = corpus.search(st['followup'], self.k, exclude=set(st['idx1']))
                ps = [corpus.passages[j][:600] for j in st['idx1'] + idx2]
                ans = clean_answer(self._call(res, p_r2(question, ps), 0.0, 32))
                step['observation'] = {'answer': ans, 'new_passages': [corpus.passages[j][:80] for j in idx2]}
                if not is_none(ans):
                    cid = C.id(ans); votes.append(('R2', cid)); raw.setdefault(cid, ans)
            step['wall_seconds'] = round(time.perf_counter() - tw, 3); res.llm_seconds += step['wall_seconds']
            res.trace.append(step)
        cl, p = self.posterior(votes)
        if cl:
            res.answer = raw[cl[int(np.argmax(p[:-1]))]]
        else:   # nothing usable observed: forced closed-book answer (charged)
            tw = time.perf_counter()
            res.answer = clean_answer(self._call(res, p_cb(question), 0.0, 24)); res.rounds += 1
            res.llm_seconds += time.perf_counter() - tw
            res.stopped_because += ' | forced closed-book answer'
        return res

    # ------------------------------------------------------------ calibration with a new LLM
    def calibrate(self, labelled: list[dict], corpus: Corpus) -> dict:
        """labelled: [{question, gold:[aliases]}]. Runs every action once per question and returns
        calibration pseudo-counts (+1 smoothing) for this LLM, retriever and corpus."""
        rel = {'R1': [1, 1], 'R2': [1, 1]}; ans = {'R1': [1, 1], 'R2': [1, 1]}
        kap = {g: [1, 1] for g in (1, 2, 3)}; pg = Counter({1: 1, 2: 1, 3: 1}); pairs = coll = 0
        dummy = Result(answer=None)
        for ex_ in labelled:
            q, gold = ex_['question'], ex_['gold'] if isinstance(ex_['gold'], list) else [ex_['gold']]
            idx = corpus.search(q, self.k); ps = [corpus.passages[j][:600] for j in idx]
            a1, fu = parse_r1(self._call(dummy, p_r1(q, ps), 0.0, 96))
            idx2 = corpus.search(fu or q, self.k, exclude=set(idx))
            a2 = clean_answer(self._call(dummy, p_r2(q, ps + [corpus.passages[j][:600] for j in idx2]), 0.0, 32))
            for t, a in (('R1', a1), ('R2', a2)):
                if is_none(a): ans[t][1] += 1
                else: ans[t][0] += 1; c = gold_match(a, gold); rel[t][0] += c; rel[t][1] += 1 - c
            with ThreadPoolExecutor(self.m) as ex:
                outs = [clean_answer(o) for o in ex.map(lambda _: self._call(dummy, p_cb(q), self.temp, 24), range(self.m))]
            good = [s for s in outs if not is_none(s)]
            if good:
                C = Clusters(); ids = [C.id(s) for s in good]; mo, g = Counter(ids).most_common(1)[0]
                ok = gold_match(good[ids.index(mo)], gold); kap[g][0] += ok; kap[g][1] += 1 - ok; pg[g] += 1
            wrong = [x for x in (a1, a2, outs[0] if outs else None) if not is_none(x) and not gold_match(x, gold)]
            C = Clusters(); ids = [C.id(x) for x in wrong]
            for i in range(len(ids)):
                for j in range(i + 1, len(ids)): pairs += 1; coll += ids[i] == ids[j]
        tot = sum(pg.values())
        return dict(rel=rel, ans=ans, kap={str(g): v for g, v in kap.items()},
                    pg={str(g): pg[g] / tot for g in pg}, eta=(1 + coll) / (2 + pairs),
                    n_questions=len(labelled), llm_calls=dummy.llm_calls)