"""BRAG - Bayesian RAG: Gradio Space (live mode) with a programmatic API."""
import json, os, hashlib, threading
from functools import lru_cache
import gradio as gr
from brag import BRAG, Corpus, hf_llm

HERE = os.path.dirname(os.path.abspath(__file__))
CORPORA = {'HotpotQA': 'hotpotqa_corpus.json', '2WikiMultihopQA': '2wikimultihopqa_corpus.json',
           'MuSiQue': 'musique_corpus.json'}
DEFAULT_MODEL = os.getenv('BRAG_MODEL', 'Qwen/Qwen2.5-7B-Instruct')
CAL_PATH = os.path.join(HERE, 'data', 'calibration.json')
DEFAULT_CAL = json.load(open(CAL_PATH)) if os.path.exists(CAL_PATH) else None
_lock = threading.Lock()

@lru_cache(maxsize=len(CORPORA))
def builtin_corpus(name):
    return Corpus.from_json(os.path.join(HERE, 'data', CORPORA[name]))

@lru_cache(maxsize=32)
def custom_corpus(digest, text):
    text = text.strip()
    if text.startswith('['):
        rows = json.loads(text)
        ps = [r if isinstance(r, str) else f"{r.get('title', '')}: {r.get('text', '')}" for r in rows]
    else:
        ps = [l.strip() for l in text.splitlines() if l.strip()]
    return Corpus(ps)

def get_corpus(corpus, documents):
    if documents and documents.strip():
        return custom_corpus(hashlib.sha1(documents.encode()).hexdigest(), documents)
    if corpus not in CORPORA:
        raise gr.Error(f'corpus must be one of {list(CORPORA)} or documents must be given')
    with _lock:
        return builtin_corpus(corpus)

def get_llm(model, hf_token, provider):
    if os.getenv('BRAG_MOCK') == '1':                       # local testing only
        from mock_llm import mock_llm
        return mock_llm
    token = (hf_token or '').strip() or os.getenv('HF_TOKEN')
    if not token:
        raise gr.Error('No Hugging Face token: pass hf_token, or set the HF_TOKEN secret on the Space.')
    return hf_llm((model or DEFAULT_MODEL).strip(), token, (provider or 'auto').strip())

def parse_cal(calibration):
    if calibration and str(calibration).strip():
        try:
            return json.loads(calibration) if isinstance(calibration, str) else calibration
        except json.JSONDecodeError as e:
            raise gr.Error(f'calibration is not valid JSON: {e}')
    return DEFAULT_CAL

def brag(question: str, corpus: str = 'HotpotQA', documents: str = '', lam: float = 0.05,
         model: str = DEFAULT_MODEL, hf_token: str = '', provider: str = 'auto', calibration: str = ''):
    """Answer a question with BRAG (Bayesian RAG).

    Args:
        question: the question to answer.
        corpus: built-in corpus to search: HotpotQA, 2WikiMultihopQA or MuSiQue (ignored if documents given).
        documents: optional custom corpus: one passage per line, or a JSON list of strings / {title, text}.
        lam: cost of one LLM call in units of answer error (higher = cheaper, earlier stopping).
        model: Hugging Face model id, or URL of an OpenAI-compatible endpoint (dedicated Inference Endpoint, TGI, vLLM).
        hf_token: optional Hugging Face token (defaults to the Space's HF_TOKEN secret).
        provider: Hugging Face inference provider ('auto' picks one that serves the model).
        calibration: optional JSON calibration (output of the calibrate endpoint) for this model.
    Returns:
        answer, and a JSON report with the per-step posterior, information rates, actions and costs.
    """
    if not question or not question.strip():
        raise gr.Error('question is empty')
    engine = BRAG(get_llm(model, hf_token, provider), parse_cal(calibration), lam=float(lam))
    try:
        res = engine.run(question.strip(), get_corpus(corpus, documents))
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f'LLM/provider error: {type(e).__name__}: {e}')
    rep = res.to_dict(); rep['model'] = model; rep['lambda'] = float(lam)
    return res.answer or '', rep

def calibrate(labelled_json: str, corpus: str = 'HotpotQA', documents: str = '',
              model: str = DEFAULT_MODEL, hf_token: str = '', provider: str = 'auto'):
    """Calibrate BRAG's hyperpriors for a model/corpus from labelled questions.

    Args:
        labelled_json: JSON list of {"question": ..., "gold": [aliases]} (30+ recommended, max 200).
        corpus / documents / model / hf_token / provider: as in the brag endpoint.
    Returns:
        calibration JSON to pass as `calibration` to the brag endpoint.
    """
    try:
        rows = json.loads(labelled_json)
        assert isinstance(rows, list) and rows and all('question' in r and 'gold' in r for r in rows)
    except Exception:
        raise gr.Error('labelled_json must be a JSON list of {"question": ..., "gold": [...]}')
    if len(rows) > 200:
        raise gr.Error('at most 200 labelled questions per call')
    engine = BRAG(get_llm(model, hf_token, provider), DEFAULT_CAL)
    try:
        return engine.calibrate(rows, get_corpus(corpus, documents))
    except Exception as e:
        raise gr.Error(f'LLM/provider error: {type(e).__name__}: {e}')

API_DOC = """
### Programmatic use
```python
from gradio_client import Client
c = Client("Angshul/BRAG")                     # or Client("Angshul/BRAG", hf_token="hf_...")
answer, report = c.predict(
    question="Friedrich Ratzel coined the phrase Lebensraum; which political movement made it notorious?",
    corpus="HotpotQA", documents="", lam=0.05,
    model="Qwen/Qwen2.5-7B-Instruct", hf_token="", provider="auto", calibration="",
    api_name="/brag")
```
Own documents: pass `documents` (one passage per line, or a JSON list). Different model: run
`api_name="/calibrate"` once on 30+ labelled questions and pass the returned JSON as `calibration`.
REST: `POST /gradio_api/call/brag` with `{"data": [question, corpus, documents, lam, model, hf_token, provider, calibration]}`.
"""

with gr.Blocks(title='BRAG - Bayesian RAG') as demo:
    gr.Markdown('# BRAG — Bayesian RAG\nClosed-form Bayesian orchestration of retrieval-augmented QA. '
                'The orchestrator makes **no LLM calls**: each decision is arithmetic on conjugate posteriors '
                '(microseconds). It chooses between reading retrieved passages, a second retrieval hop and '
                'closed-book sampling, and stops when the expected error reduction per call falls below λ.')
    with gr.Tab('Ask'):
        q = gr.Textbox(label='Question', lines=2)
        with gr.Row():
            corpus = gr.Dropdown(list(CORPORA), value='HotpotQA', label='Corpus')
            lam = gr.Slider(0.005, 0.3, value=0.05, step=0.005, label='λ (cost per LLM call)')
        docs = gr.Textbox(label='Own documents (optional): one passage per line, or JSON list', lines=3)
        with gr.Accordion('Model', open=False):
            model = gr.Textbox(value=DEFAULT_MODEL, label='HF model id or OpenAI-compatible endpoint URL')
            provider = gr.Textbox(value='auto', label='Inference provider')
            tok = gr.Textbox(label='HF token (optional; default: Space secret)', type='password')
            cal = gr.Textbox(label='Calibration JSON (optional)', lines=2)
        go = gr.Button('Answer', variant='primary')
        ans = gr.Textbox(label='Answer'); rep = gr.JSON(label='Trace: posterior, information rates, actions, costs')
        go.click(brag, [q, corpus, docs, lam, model, tok, provider, cal], [ans, rep], api_name='brag')
    with gr.Tab('Calibrate'):
        lab = gr.Textbox(label='Labelled questions JSON: [{"question": ..., "gold": [...]}, ...]', lines=6)
        with gr.Row():
            corpus2 = gr.Dropdown(list(CORPORA), value='HotpotQA', label='Corpus')
            model2 = gr.Textbox(value=DEFAULT_MODEL, label='HF model id')
        docs2 = gr.Textbox(label='Own documents (optional)', lines=2)
        tok2 = gr.Textbox(label='HF token (optional)', type='password'); prov2 = gr.Textbox(value='auto', label='Provider')
        go2 = gr.Button('Calibrate'); out2 = gr.JSON(label='Calibration (pass to the brag endpoint)')
        go2.click(calibrate, [lab, corpus2, docs2, model2, tok2, prov2], out2, api_name='calibrate')
    gr.Markdown(API_DOC)

demo.queue(default_concurrency_limit=int(os.getenv('BRAG_CONCURRENCY', '16')), max_size=256)

if __name__ == '__main__':
    demo.launch()