"""Disposable feasibility spike: multilingual retrieval + grounded answers with a small local model.

Usage: python3 spike.py embed            # embed library + questions with bge-m3, then unload it
       python3 spike.py run <model-tag>  # answer all questions with one LLM, then unload it
Without SPIKE_SET it runs test 1 (questions.json: English, Dutch, French, 22 Oct 2026). Set SPIKE_SET=lang2 for
test 2 (questions_lang2.json: German, Spanish, Polish, Swedish, 25 Oct 2026; results get a _lang2 suffix).
Published with the Waystone survival AI benchmark: https://waystoneoffline.com/guides/survival-ai-benchmark/
Only uses the local Ollama server (http://localhost:11434) and the Python standard library.
"""
import json, math, os, re, subprocess, sys, threading, time, urllib.request
from pathlib import Path

HERE = Path(__file__).parent
OUT = HERE / "results"
OUT.mkdir(exist_ok=True)
OLLAMA = "http://localhost:11434"
EMBED_MODEL = "bge-m3"
TOP_K = 3
NUM_CTX = 4096

LIB = json.loads((HERE / "library.json").read_text())
SET = os.environ.get("SPIKE_SET", "")
SUFFIX = f"_{SET}" if SET else ""
QS = json.loads((HERE / f"questions{SUFFIX}.json").read_text())


def post(path, body, stream=False):
    req = urllib.request.Request(OLLAMA + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
    resp = urllib.request.urlopen(req, timeout=600)
    if stream:
        return resp
    return json.loads(resp.read())


def unload(model):
    if model == EMBED_MODEL:
        post("/api/embed", {"model": model, "input": "x", "keep_alive": 0})
    else:
        post("/api/generate", {"model": model, "keep_alive": 0})


class RamSampler(threading.Thread):
    """Samples total resident memory of all ollama processes every 0.5 s."""
    def __init__(self):
        super().__init__(daemon=True)
        self.peak_mb, self.stop = 0.0, False

    def run(self):
        while not self.stop:
            out = subprocess.run(["ps", "-axo", "rss=,args="], capture_output=True, text=True).stdout
            total = sum(int(l.split()[0]) for l in out.splitlines() if "Ollama.app" in l or "ollama" in l.lower())
            self.peak_mb = max(self.peak_mb, total / 1024)
            time.sleep(0.5)


def cos(a, b):
    return sum(x * y for x, y in zip(a, b)) / (math.sqrt(sum(x * x for x in a)) * math.sqrt(sum(y * y for y in b)))


def all_questions():
    qs = [dict(q) for q in QS["single"]]
    for t in QS["conversation"]["turns"]:
        qs.append(dict(t, country=QS["conversation"]["country"]))
    return qs


def cmd_embed():
    ram = RamSampler(); ram.start()
    t0 = time.time()
    post("/api/embed", {"model": EMBED_MODEL, "input": "warm up"})
    load_s = time.time() - t0
    # Passages are embedded once, in English only: the library is not translated.
    pv = post("/api/embed", {"model": EMBED_MODEL, "input": [p["title"] + ". " + p["text"] for p in LIB]})["embeddings"]
    qv, qtimes = {}, []
    for q in all_questions():
        t = time.time()
        qv[q["id"]] = post("/api/embed", {"model": EMBED_MODEL, "input": q["q"]})["embeddings"][0]
        qtimes.append(time.time() - t)
    unload(EMBED_MODEL)
    ram.stop = True
    retrieval = {}
    for q in all_questions():
        allowed = [i for i, p in enumerate(LIB) if p["country"] in ("all", q["country"])]
        scored = sorted(((cos(qv[q["id"]], pv[i]), LIB[i]["id"]) for i in allowed), reverse=True)
        retrieval[q["id"]] = [{"id": pid, "score": round(s, 4)} for s, pid in scored[:TOP_K]]
    meta = {"embed_model": EMBED_MODEL, "load_s": round(load_s, 2), "peak_ram_mb": round(ram.peak_mb),
            "avg_query_embed_ms": round(1000 * sum(qtimes) / len(qtimes), 1)}
    (OUT / f"retrieval{SUFFIX}.json").write_text(json.dumps({"meta": meta, "retrieval": retrieval}, indent=1, ensure_ascii=False))
    print(json.dumps(meta))


SYSTEM = """You are an offline emergency-preparedness assistant. You may ONLY use the numbered library passages given below.
Rules:
1. Answer in the SAME language as the user's latest message (Dutch, English or French), even if the passages are in English and even if earlier messages used another language.
2. After every fact, cite its passage id in square brackets, e.g. [P02].
3. Copy numbers and units exactly as the passage gives them (you may add a translated unit word). Never change a number.
4. Keep warnings and prohibitions exact: if a passage says NOT or NEVER, your answer must say so clearly.
5. If the passages do not answer the question, reply with the single word NOT_IN_LIBRARY followed by one short sentence, in the user's language, saying the library does not cover this. Do not use general knowledge.
6. Be short: at most 4 sentences."""
if SET == "ro":
    SYSTEM = SYSTEM.replace("(Dutch, English or French)", "(Romanian, Dutch, English or French)")
elif SET == "lang2":
    SYSTEM = SYSTEM.replace("(Dutch, English or French)", "(German, Spanish, Polish, Swedish, Dutch, English or French)")


def passages_block(ids):
    by = {p["id"]: p for p in LIB}
    return "\n".join(f"[{i}] {by[i]['title']}: {by[i]['text']}" for i in ids)


def chat(model, messages):
    t0 = time.time()
    ttft, text, final = None, "", {}
    resp = post("/api/chat", {"model": model, "messages": messages, "stream": True,
                              "options": {"temperature": 0, "num_ctx": NUM_CTX, "seed": 1}}, stream=True)
    for line in resp:
        ch = json.loads(line)
        piece = ch.get("message", {}).get("content", "")
        if piece and ttft is None:
            ttft = time.time() - t0
        text += piece
        if ch.get("done"):
            final = ch
    return text, {"ttft_s": round(ttft or 0, 2), "total_s": round(time.time() - t0, 2),
                  "eval_tokens": final.get("eval_count"),
                  "tok_per_s": round(final["eval_count"] / (final["eval_duration"] / 1e9), 1) if final.get("eval_duration") else None,
                  "prompt_tokens": final.get("prompt_eval_count"),
                  "prompt_tok_per_s": round(final["prompt_eval_count"] / (final["prompt_eval_duration"] / 1e9), 1) if final.get("prompt_eval_duration") else None}


STOP = {
    "nl": "de het een en van is niet je u op te dat voor met zijn wordt ook geen moet water dan als bij".split(),
    "en": "the a an and of is not you on to that for with are be also no should water then if at".split(),
    "ro": "și în este nu un o de la cu pe pentru că se să mai sau din apa trebuie dacă sunt care".split(),
    "fr": "le la les un une et de est pas vous sur pour que avec sont être aussi ne doit eau si au des du".split(),
    "de": "der die das und ist nicht ein eine zu den mit sie auf für von im es wasser sollten kein nie oder auch".split(),
    "es": "el los las y es no que con por del al más como su agua debe una lo hay para pero".split(),
    "pl": "i w nie na się z do to że jest jak lub przez dla po czy woda wody od tylko należy".split(),
    "sv": "och att det som på är inte för med av till har om vatten kan ska du eller vid ett".split(),
}


def detect_lang(text):
    words = re.findall(r"[a-zàâçéèêëîïôûùüÿœăâîșțşţäößåñáíóúąćęłńśźż']+", text.lower())
    scores = {l: sum(w in s for w in words) for l, s in STOP.items()}
    return max(scores, key=scores.get)


def nums(s):
    s = re.sub(r"(\d)[   ,](\d{3})\b", r"\1\2", s)
    return set(re.findall(r"\d+(?:[.,]\d+)?", s.replace(",", ".")))


def score(q, retrieved, answer):
    ids = [r["id"] for r in retrieved]
    refused = "NOT_IN_LIBRARY" in answer
    cited = re.findall(r"\[([A-Z]{1,2}\d{1,2})\]", answer)
    res = {"refused": refused, "cited": sorted(set(cited)), "answer_lang": detect_lang(answer.replace("NOT_IN_LIBRARY", ""))}
    res["lang_ok"] = res["answer_lang"] == q["lang"]
    if q["expect"] is None:
        res["refusal_ok"] = refused
        return res
    res["refusal_ok"] = not refused
    res["hit_top1"] = ids[0] in q["expect"]
    res["hit_top3"] = any(i in q["expect"] for i in ids)
    res["citation_ok"] = bool(cited) and set(cited) <= set(ids) and any(c in q["expect"] for c in cited)
    res["must_ok"] = all(re.search(m, answer, re.I) for m in q["must"])
    res["neg_ok"] = True if not q["neg"] else bool(re.search(q["neg"], answer, re.I))
    allowed = nums(passages_block([c for c in set(cited) if c in ids] or ids)) | nums(q["q"])
    clean = re.sub(r"\[[A-Z]{1,2}\d{1,2}\]", "", answer)
    res["unsupported_numbers"] = sorted(nums(clean) - allowed)
    return res


def cmd_run(model):
    R = json.loads((OUT / f"retrieval{SUFFIX}.json").read_text())["retrieval"]
    ram = RamSampler(); ram.start()
    t0 = time.time()
    post("/api/generate", {"model": model, "prompt": "", "options": {"num_ctx": NUM_CTX}})
    load_s = time.time() - t0
    rows = []
    for q in QS["single"]:
        ret = R[q["id"]]
        msgs = [{"role": "system", "content": SYSTEM + "\n\nLibrary passages:\n" + passages_block([r["id"] for r in ret])},
                {"role": "user", "content": q["q"]}]
        ans, perf = chat(model, msgs)
        rows.append({"id": q["id"], "lang": q["lang"], "q": q["q"], "retrieved": ret, "answer": ans, "perf": perf, "score": score(q, ret, ans)})
        print(q["id"], perf["total_s"], "s", json.dumps(rows[-1]["score"], ensure_ascii=False), flush=True)
    # Multi-turn conversation that switches language each turn; passages are refreshed for each turn.
    hist = []
    for t in QS["conversation"]["turns"]:
        ret = R[t["id"]]
        msgs = [{"role": "system", "content": SYSTEM + "\n\nLibrary passages:\n" + passages_block([r["id"] for r in ret])}] + hist + [{"role": "user", "content": t["q"]}]
        ans, perf = chat(model, msgs)
        hist += [{"role": "user", "content": t["q"]}, {"role": "assistant", "content": ans}]
        rows.append({"id": t["id"], "lang": t["lang"], "q": t["q"], "retrieved": ret, "answer": ans, "perf": perf, "score": score(t, ret, ans)})
        print(t["id"], perf["total_s"], "s", json.dumps(rows[-1]["score"], ensure_ascii=False), flush=True)
    ps = json.loads(urllib.request.urlopen(OLLAMA + "/api/ps").read())
    unload(model)
    time.sleep(2)
    ram.stop = True
    meta = {"model": model, "load_s": round(load_s, 2), "peak_ram_mb_all_ollama": round(ram.peak_mb),
            "ollama_ps": ps.get("models")}
    safe = model.replace(":", "_").replace("/", "_")
    (OUT / f"answers_{safe}{SUFFIX}.json").write_text(json.dumps({"meta": meta, "rows": rows}, indent=1, ensure_ascii=False))
    print(json.dumps(meta, default=str))


if __name__ == "__main__":
    if sys.argv[1] == "embed":
        cmd_embed()
    else:
        cmd_run(sys.argv[2])