"""Same hamster test, but through the new Session.chat() (which uses eval+sample)."""
import sys, argparse
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
from memba import Session

MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE_DIR = "/tmp/diag_session_nemotron"
SESSION   = "hamster"


def build():
    s = Session(model_path=MODEL, session_id=SESSION, state_dir=STATE_DIR,
                n_gpu_layers=-1, n_ctx=4096, chat_format="chatml")
    # Wipe any previous state file so it's a real fresh run
    p = Path(STATE_DIR) / f"{SESSION}.memb"
    if p.exists(): p.unlink()
    print(f"[build] ack: {s.chat('My pet hamster is named Bartholomew. He is 4 years old. Reply noted.', max_tokens=40)!r}")
    print(f"[build] in-process query: {s.chat('What is the name of my pet?', max_tokens=80)!r}")
    s.save()
    print(f"[build] saved")


def query():
    s = Session(model_path=MODEL, session_id=SESSION, state_dir=STATE_DIR,
                n_gpu_layers=-1, n_ctx=4096, chat_format="chatml")
    print(f"[query] loaded state, n_tokens={s._llama.n_tokens}")
    print(f"[query] CROSS-PROCESS query: {s.chat('What is the name of my pet?', max_tokens=80)!r}")


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