"""
sim_multilingual.py — measure answer quality across languages.

Takes a sample of real KB questions, translates each into several languages, runs them through
the assistant, and scores per language:
  - answered vs wrongly escalated (should all be answerable)
  - expected source retrieved (did cross-lingual retrieval find the right English entry)
  - avg best similarity
  - reply in the right language (judge)
  - faithful to the English source answer (judge)

This tells us whether the current multilingual approach is robust or whether we need to translate
the query to English before retrieval.

Run:  python3 eval/sim_multilingual.py
"""
import json
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

import app  # noqa: E402
from app import ChatRequest, retrieve, retrieval_query, client  # noqa: E402
from google.genai import types  # noqa: E402

MODEL = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")
LANGS = ["Spanish", "French", "Hinglish (romanized Hindi)", "Tamil", "Arabic"]
SAMPLE = 50
WORKERS = 6

KB = [e for e in json.load(open(os.path.join(ROOT, "data", "kb.json")))
      if not e.get("coach_only") and not e.get("has_placeholder")]
SAMPLE_KB = KB[::max(1, len(KB) // SAMPLE)][:SAMPLE]

ARR = {"type": "array", "items": {"type": "string"}}
JUDGE = {"type": "object", "properties": {"right_language": {"type": "boolean"},
                                          "faithful": {"type": "boolean"}},
         "required": ["right_language", "faithful"]}


def _arr(prompt, n):
    for a in range(3):
        try:
            r = client.models.generate_content(model=MODEL, contents=prompt,
                config=types.GenerateContentConfig(response_mime_type="application/json",
                                                   response_schema=ARR, temperature=0.3))
            arr = json.loads(r.text)
            if isinstance(arr, list) and len(arr) == n:
                return arr
        except Exception:
            time.sleep(2 * (a + 1))
    return None


def translate(questions, lang):
    numbered = "\n".join(f"{i+1}. {q}" for i, q in enumerate(questions))
    return _arr(f"Translate each question into {lang}, the way a real user would type it. "
                f"Return a JSON array of {len(questions)} translations, same order.\n\n{numbered}", len(questions))


def judge(lang, question, reply, english_answer):
    prompt = (f"The user wrote in {lang}.\nUSER: {question}\nREPLY: {reply}\n\n"
              f"APPROVED ENGLISH ANSWER: {english_answer}\n\n"
              f"right_language = is the REPLY written in {lang}? "
              f"faithful = does the REPLY convey the same facts as the approved English answer without adding anything?")
    try:
        r = client.models.generate_content(model=MODEL, contents=prompt,
            config=types.GenerateContentConfig(response_mime_type="application/json",
                                               response_schema=JUDGE, temperature=0.0))
        return json.loads(r.text)
    except Exception:
        return None


def run_lang(lang):
    trans = translate([e["question"] for e in SAMPLE_KB], lang) or [e["question"] for e in SAMPLE_KB]
    rows = []
    for e, q in zip(SAMPLE_KB, trans):
        try:
            hits, sim = retrieve(retrieval_query(ChatRequest(message=q)))
            out = app.chat(ChatRequest(message=q))
            answered = not out["escalate"] and out.get("escalate_reason") != "out_of_scope"
            j = judge(lang, q, out["answer"], e["answer"]) if answered else None
            rows.append({"q": q, "sim": sim, "answered": answered,
                         "retrieved": e["id"] in [h["id"] for h, _ in hits],
                         "right_language": (j or {}).get("right_language"),
                         "faithful": (j or {}).get("faithful")})
        except Exception:
            rows.append({"q": q, "sim": 0, "answered": False, "retrieved": False,
                         "right_language": None, "faithful": None})
    return lang, rows


def pct(a, b):
    return f"{(100*a)//max(b,1)}%"


def main():
    print(f"Testing {len(SAMPLE_KB)} questions x {len(LANGS)} languages...")
    with ThreadPoolExecutor(max_workers=len(LANGS)) as ex:
        results = dict(ex.map(run_lang, LANGS))

    L = ["# Multilingual quality report\n", f"{len(SAMPLE_KB)} sampled questions per language\n"]
    L.append("| language | answered | source retrieved | avg sim | right language | faithful |")
    L.append("|---|---|---|---|---|---|")
    all_rows = []
    for lang in LANGS:
        rows = results[lang]
        all_rows += rows
        ans = [r for r in rows if r["answered"]]
        retr = [r for r in rows if r["retrieved"]]
        avg = sum(r["sim"] for r in rows) / max(len(rows), 1)
        rl = [r for r in ans if r["right_language"]]
        fa = [r for r in ans if r["faithful"]]
        L.append(f"| {lang} | {pct(len(ans),len(rows))} | {pct(len(retr),len(rows))} | "
                 f"{avg:.3f} | {pct(len(rl),len(ans))} | {pct(len(fa),len(ans))} |")

    # weakest cases (escalated when should answer, or low sim)
    weak = sorted([r for rows in results.values() for r in rows if not r["answered"] or r["sim"] < 0.6],
                  key=lambda r: r["sim"])[:15]
    L.append("\n## Weakest cases (retrieval risk)\n")
    for r in weak:
        L.append(f"- sim={r['sim']:.3f} answered={r['answered']} | {r['q'][:60]}")

    open(os.path.join(ROOT, "eval", "sim_multilingual_report.md"), "w").write("\n".join(L))
    json.dump({k: v for k, v in results.items()},
              open(os.path.join(ROOT, "eval", "sim_multilingual_results.json"), "w"), ensure_ascii=False, indent=2)
    print("\n".join(L))
    print("\nReport -> eval/sim_multilingual_report.md")


if __name__ == "__main__":
    main()