"""
Diagnostic: does the model retain context AFTER save in the SAME process,
and AFTER load in a fresh process? Compares three scenarios.
"""
import sys, argparse
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
from llama_cpp import Llama
from memba import core

MODEL = "/home/emil/Desktop/Coding/AI/Memba/falcon-mamba-7B-instruct-Q4_K_M.gguf"
STATE = "/tmp/diag_saveload.memb"

INGEST  = "I'm telling you a secret. My pet hamster is named Bartholomew. He is 4 years old. Reply 'ok'."
QUERY   = "What is the name of my pet?"


def chatml(msg):
    return f"<|im_start|>user\n{msg}<|im_end|>\n<|im_start|>assistant\n"


def make_llama():
    return Llama(model_path=MODEL, n_ctx=2048, n_gpu_layers=-1, verbose=False)


def ask(m, prompt):
    out = m(chatml(prompt), max_tokens=60, stop=["<|im_end|>"], echo=False)
    return out["choices"][0]["text"].strip()


def build():
    m = make_llama()
    ack = ask(m, INGEST)
    print(f"  [build] ack: {ack!r}")
    print(f"  [build] state size (live): {core.get_state_size(m):,} B")
    # In-process query BEFORE saving
    print(f"  [build] in-proc query BEFORE save: {ask(m, QUERY)!r}")
    # Save
    core.save_state(m, MODEL, STATE)
    print(f"  [build] state saved")
    # In-process query AFTER saving (should still work — save shouldn't mutate)
    print(f"  [build] in-proc query AFTER save:  {ask(m, QUERY)!r}")


def query():
    m = make_llama()
    print(f"  [query] before load — fresh model: {ask(m, QUERY)!r}")
    core.load_state(m, MODEL, STATE)
    print(f"  [query] state size after load: {core.get_state_size(m):,} B")
    print(f"  [query] after load:                {ask(m, QUERY)!r}")
    # Try a second time in case position is wonky
    print(f"  [query] second ask:                {ask(m, QUERY)!r}")


if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else "build"
    {"build": build, "query": query}[cmd]()