"""Grounded question answering: retrieve passages, then let Gemini answer ONLY from them, with citations.

Setup: put the key for the provider you use in FinalProject.env (git-ignored):
    GEMINI_API_KEY=...        pip install google-genai        --provider gemini    (default)
    ANTHROPIC_API_KEY=...     pip install anthropic           --provider anthropic
    OPENAI_API_KEY=...        pip install openai              --provider openai

    python rag.py "ကျေးလက်ဒေသ ဆေးရုံများတွင် ဆရာဝန် မလုံလောက်မှုကို အစိုးရက ဘယ်လိုဖြေရှင်းမလဲ"
    python rag.py --provider anthropic --batch data/pilot_questions.jsonl --out data/pilot_answers.jsonl
    python rag.py --no-retrieval "..."       # ablation: same model, no passages (for the faithfulness comparison)
    python rag.py --embed sealion --mode hybrid --per-speech 2 --batch data/eval_questions.jsonl --out data/eval_answers.jsonl
    python rag.py --provider anthropic --list-models / --probe

Retrieval configuration: --embed <model key> (embed_models.py), --mode dense|bm25|hybrid, --k 5, --per-speech N
(at most N passages from the same speech — the pilot's top-k was filled by one long speech).
"""
import argparse, json, logging, os, sys, time
from dotenv import load_dotenv
logging.getLogger("google_genai").setLevel(logging.ERROR)   # silence the harmless AFC warning
from common import citation_of
from retrieval import Retriever

PROVIDERS = {   # provider: (env var, default model, fallback models tried in order when the default is overloaded)
    # probe 30 Sep 2026: 3.8/3.7/3.5-flash all 503 "high demand" on the free tier; flash-latest, 3.5-flash-lite and gemma-4 answered
    "gemini":    ("GEMINI_API_KEY",    "gemini-3.5-flash-lite", ["gemma-4-26b-a4b-it", "gemini-flash-latest", "gemini-3.8-flash"]),
    "anthropic": ("ANTHROPIC_API_KEY", "claude-haiku-5-5", []),   # grader: Haiku ≈ 1 cent per answer (Sonnet was ≈ 6 cents)      # run --provider anthropic --list-models to see current names
    "openai":    ("OPENAI_API_KEY",    "gpt-5-mini",        []),
}
DEFAULT_PROVIDER = "gemini"
DEFAULT_MODEL = None    # None = the provider's default
K = 5
EMBED = "sealion"; MODE = "hybrid"; PER_SPEECH = 2      # SQ1 winner (2 Nov 2026); BGE-M3 was the pilot model

SYSTEM = """You are an assistant that explains what was said in Myanmar's parliament (Hluttaw) between 2011 and 2020.
You will be given a question and numbered passages. The passages are official parliamentary transcripts written in
Burmese; each begins with a header giving the house, date, agenda section, item and speaker.
Rules:
1. Answer ONLY from the passages. Do not add outside knowledge or general background.
2. Answer in the language of the QUESTION: a Burmese question gets a Burmese answer; an English question gets an
   English answer that reports what the Burmese passages say (translate the content; keep names as in the passages).
   Never refuse or switch language because the passages are in Burmese — read them and report them.
3. Write 3–8 plain sentences a non-expert can follow. End every sentence that states a fact with the number(s) of
   the passage(s) it comes from, like [2] or [1][3]. A sentence without a citation must not contain a fact.
4. Say who said it (the MP, minister or ministry, as named in the passage header) and on what date.
5. Copy numbers, dates, place names and amounts exactly as they appear in the passages. Do not round, convert or
   infer numbers that are not there.
6. If the passages do not answer the question, say so in one sentence (in the question's language) and do not guess.
   If they answer only part of it, answer that part and say what is missing.
7. If the question asks for an overview of what a person, party or house discussed "most" or "overall", say in one
   sentence that the answer rests only on the retrieved passages, then report what those passages show.
8. Speaker headers give each MP's constituency and, where known, party (e.g. "· အမျိုးသားဒီမိုကရေစီအဖွဲ့ချုပ် (NLD)").
   For a question about a party or group, use those labels to attribute statements; if no passage carries that label,
   say the retrieved record does not identify speakers from that party, then still report what the passages say on the topic."""

LANG_RULE = {"my": "ANSWER LANGUAGE: Burmese. Write the whole answer in Burmese, whatever language the question is in.",
             "en": "ANSWER LANGUAGE: English. Write the whole answer in English, reporting what the Burmese passages say "
                   "(translate the content; keep personal names as in the passages, with a Latin transliteration in brackets if helpful)."}

def build_prompt(question, hits, lang=None):
    ctx = "\n\n".join(f"[{i}] {citation_of(m)}\n{m['text']}" for i, (s, m) in enumerate(hits, 1))
    rule = ("\n" + LANG_RULE[lang]) if lang in LANG_RULE else ""        # None / "auto": rule 2 of SYSTEM applies
    return f"{SYSTEM}{rule}\n\nPASSAGES:\n{ctx}\n\nQUESTION: {question}\n\nANSWER:"

def build_prompt_no_retrieval(question):
    return (f"Answer the following question about what was said in Myanmar's parliament (Hluttaw) between 2011 and 2020. "
            f"Answer in the language of the question, in 3–8 plain sentences, naming who said what and when as precisely as you can.\n\n"
            f"QUESTION: {question}\n\nANSWER:")

class Generator:
    def __init__(self, provider=DEFAULT_PROVIDER, model=DEFAULT_MODEL, no_fallback=False):
        load_dotenv("FinalProject.env")
        env, default, fallbacks = PROVIDERS[provider]
        if no_fallback: fallbacks = []
        key = os.environ.get(env)
        if not key: sys.exit(f"{env} missing — add it to FinalProject.env")
        self.provider, self.model = provider, (model or default)
        self.fallbacks = [m for m in fallbacks if m != self.model]
        if provider == "gemini":
            from google import genai; self.client = genai.Client(api_key=key)
        elif provider == "anthropic":
            import anthropic; self.client = anthropic.Anthropic(api_key=key)
        elif provider == "openai":
            import openai; self.client = openai.OpenAI(api_key=key)

    def _generate(self, model, prompt):
        """One call to the provider; returns the answer text."""
        if self.provider == "gemini":
            return self.client.models.generate_content(model=model, contents=prompt).text or ""
        if self.provider == "anthropic":
            r = self.client.messages.create(model=model, max_tokens=1500, messages=[{"role": "user", "content": prompt}])
            return "".join(b.text for b in r.content if getattr(b, "type", "") == "text")
        if self.provider == "openai":
            r = self.client.chat.completions.create(model=model, messages=[{"role": "user", "content": prompt}])
            return r.choices[0].message.content or ""

    def probe(self, candidates):
        """Try each model once with a tiny prompt; report which ones respond."""
        ok = []
        for m in candidates:
            t = time.time()
            try:
                text = self._generate(m, "Reply with the single word OK.")
                print(f"  {m:32} OK   ({time.time()-t:.1f}s)  {text.strip()[:40]!r}"); ok.append(m)
            except Exception as e:
                msg = str(e); code = getattr(e, "code", None) or getattr(e, "status_code", None)
                short = "overloaded (503)" if "UNAVAILABLE" in msg else "rate-limited (429)" if "RESOURCE_EXHAUSTED" in msg else f"error {code}: {msg[:600]}"
                print(f"  {m:32} FAIL {short}")
        print("\nworking:", ", ".join(ok) if ok else "none — wait a while and rerun, or try another --provider")

    def list_models(self):
        print(f"models available to this {self.provider} key:")
        if self.provider == "gemini":
            for m in self.client.models.list():
                if "generateContent" in (getattr(m, "supported_actions", None) or []): print("  ", m.name.replace("models/", ""))
        elif self.provider == "anthropic":
            for m in self.client.models.list(limit=100): print("  ", m.id)
        elif self.provider == "openai":
            for m in sorted(x.id for x in self.client.models.list()): print("  ", m)

    def __call__(self, prompt):
        """Returns (text, model_used). Retries on rate limits / overload; after 3 overload failures switches to the next fallback model."""
        while True:
            for attempt in range(3):
                try:
                    return self._generate(self.model, prompt), self.model
                except Exception as e:
                    msg = str(e)[:300]
                    code = getattr(e, "code", None) or getattr(e, "status_code", None)
                    retryable = code in (429, 500, 502, 503, 504, 529) or "RESOURCE_EXHAUSTED" in msg or "UNAVAILABLE" in msg or "overloaded" in msg.lower()
                    if "exceeded your current quota" in msg or "quota" in msg.lower() and code == 429:
                        print(f"  daily quota of {self.model} used up — switching model or stop and rerun with --resume after the reset (midnight Pacific = 2 pm Thailand)", file=sys.stderr)
                        break
                    print(f"  [{self.model}] error: {msg}", file=sys.stderr)
                    if not retryable:
                        if code == 404 or "NOT_FOUND" in msg:
                            print(f"  model '{self.model}' not found for this key — run  python rag.py --list-models  and pass --model <name>", file=sys.stderr)
                        sys.exit(1)
                    wait = 10 * (attempt + 1); print(f"  retry in {wait}s", file=sys.stderr); time.sleep(wait)
            if not self.fallbacks:
                return "[generation failed]", self.model
            nxt = self.fallbacks.pop(0)
            print(f"  {self.model} still unavailable — switching to {nxt}", file=sys.stderr)
            self.model = nxt

def is_burmese(text):
    return any("\u1000" <= ch <= "\u109f" for ch in text)

def translate_to_burmese(question, gen):
    """Query translation for cross-lingual retrieval: a non-Burmese question is rendered in Burmese so that the
    keyword index (and the dense model, in-language) can match names, places and terms in the record."""
    try:
        text, _ = gen(f"Translate the following question into natural Burmese as a Myanmar citizen would ask it about "
                      f"parliament. Keep personal names in their usual Burmese form (e.g. Aung San Suu Kyi → ဒေါ်အောင်ဆန်းစုကြည်). "
                      f"Output only the Burmese translation, nothing else.\n\n{question}")
        text = text.strip().splitlines()[0].strip() if text.strip() else ""
        return text if is_burmese(text) else ""
    except Exception:
        return ""

def answer(question, ret, gen, k=K, no_retrieval=False, per_speech=PER_SPEECH, lang=None, translate=True):
    """lang: None/"auto" = language of the question (detected, then stated explicitly to the model);
    "my" or "en" = force the answer language. translate: for non-Burmese questions, also retrieve with a Burmese translation."""
    if lang in (None, "auto"): lang = "my" if is_burmese(question) else "en"
    query_my = ""
    if not no_retrieval and translate and not is_burmese(question):
        query_my = translate_to_burmese(question, gen)
    hits = [] if no_retrieval else ret.search(question, k=k, per_speech=per_speech, extra_queries=[query_my] if query_my else ())
    prompt = build_prompt_no_retrieval(question) if no_retrieval else build_prompt(question, hits, lang)
    t = time.time(); text, used = gen(prompt); dt = time.time() - t
    return {"question": question, "answer": text.strip(), "provider": gen.provider, "model": used, "retrieval": not no_retrieval,
            "retriever": None if no_retrieval else ret.name, "k": k, "per_speech": per_speech, "lang": lang, "query_my": query_my or None, "gen_seconds": round(dt, 2),
            "passages": [{"n": i, "score": round(s, 4), "chunk_id": m["chunk_id"], "speech_id": m["speech_id"], "citation": citation_of(m), "text": m["text"]}
                         for i, (s, m) in enumerate(hits, 1)]}

def show(res):
    print(f"\nQ: {res['question']}\n\nA: {res['answer']}\n")
    for p in res["passages"]: print(f"[{p['n']}] ({p['score']:.3f}) {p['citation']}")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("question", nargs="?")
    ap.add_argument("--batch", help="jsonl with a 'question' field per line")
    ap.add_argument("--out", default=None)
    ap.add_argument("--k", type=int, default=K)
    ap.add_argument("--provider", default=DEFAULT_PROVIDER, choices=list(PROVIDERS))
    ap.add_argument("--model", default=DEFAULT_MODEL, help="model name for the provider (default: provider's default)")
    ap.add_argument("--no-retrieval", action="store_true")
    ap.add_argument("--resume", action="store_true", help="continue a batch: keep answers already in --out, skip their ids")
    ap.add_argument("--no-fallback", action="store_true", help="never switch model; stop the batch instead (use for the SQ2 runs)")
    ap.add_argument("--embed", default=EMBED, help="embedding model key (embed_models.py)")
    ap.add_argument("--mode", default=MODE, choices=["dense", "bm25", "hybrid"])
    ap.add_argument("--per-speech", type=int, default=PER_SPEECH, help="max passages per speech in the top-k (0 = no cap)")
    ap.add_argument("--list-models", action="store_true", help="show the models this API key can call, then exit")
    ap.add_argument("--probe", action="store_true", help="try each candidate model once and report which respond, then exit")
    args = ap.parse_args()
    gen = Generator(args.provider, args.model, args.no_fallback)
    if args.list_models: gen.list_models(); return
    if args.probe:
        extra = ["gemini-3.7-flash", "gemini-3.5-flash"] if args.provider == "gemini" else []
        gen.probe([gen.model] + gen.fallbacks + extra); return
    idx = None if args.no_retrieval else Retriever(model=args.embed, mode=args.mode)
    ps = args.per_speech or None
    if args.batch:
        qs = [json.loads(l) for l in open(args.batch, encoding="utf-8") if l.strip()]
        out_path = args.out or args.batch.replace(".jsonl", "_answers.jsonl")
        done = set()
        if args.resume and os.path.exists(out_path):
            done = {json.loads(l)["id"] for l in open(out_path, encoding="utf-8") if l.strip() and "[generation failed]" not in l}
            print(f"resuming: {len(done)} answers already in {out_path}")
        out = open(out_path, "a" if args.resume else "w", encoding="utf-8")
        for i, q in enumerate(qs, 1):
            qid = q.get("id", i)
            if qid in done: continue
            res = answer(q["question"], idx, gen, args.k, args.no_retrieval, ps); res["id"] = qid
            if res["answer"] == "[generation failed]":
                print(f"{i}/{len(qs)} {qid}: generation failed — stopping; rerun with --resume later", flush=True); break
            out.write(json.dumps(res, ensure_ascii=False) + "\n"); out.flush()
            print(f"{i}/{len(qs)} done ({res['gen_seconds']}s)", flush=True)
        print(f"→ {out_path}")
    elif args.question:
        show(answer(args.question, idx, gen, args.k, args.no_retrieval, ps))
    else:
        print("Ask a question (empty line to quit)")
        while True:
            try: q = input("\n> ").strip()
            except (EOFError, KeyboardInterrupt): break
            if not q: break
            show(answer(q, idx, gen, args.k, args.no_retrieval, ps))

if __name__ == "__main__":
    main()