"""
diag_nemotron2.py — same test, but in cross-process query phase we use
raw __call__ with reset=False to NOT clear KV-cache before generation.

Hypothesis: create_chat_completion() resets KV-cache on each call, which
defeats memba's loaded state. Bypassing that should let state work.
"""
import sys, time
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/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE = "/tmp/diag_nemotron.memb"   # reuse state from previous test

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


def query_raw():
    """Cross-process query, bypassing create_chat_completion's reset."""
    if not Path(STATE).exists():
        print("[query] no state — run diag_nemotron.py build first")
        return
    print(f"[query] loading state ({Path(STATE).stat().st_size:,} B)…")
    t0 = time.time()
    m = make_llama()
    print(f"[query] model loaded in {time.time()-t0:.1f}s")

    core.load_state(m, MODEL, STATE)
    print(f"[query] state loaded in {time.time()-t0:.1f}s total")

    # Continue the conversation: just append a new user turn manually.
    # Nemotron uses ChatML-like markers based on what we've seen:
    # <|im_start|>user\n…<|im_end|>\n<|im_start|>assistant\n
    # If wrong, model will still answer something interpretable.
    continuation = "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n"

    # Tokenize continuation WITHOUT BOS — we're mid-sequence
    tokens = m.tokenize(continuation.encode("utf-8"), add_bos=False, special=True)
    print(f"[query] continuation tokenized to {len(tokens)} tokens")

    # Use generate() which feeds via eval (does NOT reset KV-cache)
    print(f"[query] generating tokens (preserving loaded state)…")
    out_tokens = []
    eos = m.token_eos()
    for token in m.generate(tokens, top_k=1, temp=0.0):  # greedy
        if token == eos:
            break
        out_tokens.append(token)
        if len(out_tokens) >= 200:
            break
        # Also stop at <|im_end|>
        snippet = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
        if "<|im_end|>" in snippet:
            break

    text = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
    print(f"\n[query] answer:\n{text}")


def query_reset_default():
    """Control: same continuation but with default reset=True (clears KV)."""
    if not Path(STATE).exists():
        return
    print(f"\n[control] same continuation, default reset=True (should fail):")
    m = make_llama()
    core.load_state(m, MODEL, STATE)
    out = m(
        "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n",
        max_tokens=120, temperature=0.1, echo=False,  # reset=True is default
    )
    print(out["choices"][0]["text"])


if __name__ == "__main__":
    query_raw()
    # query_reset_default()  # optionally enable for direct comparison