"""
01_basic_save_load.py — save and load SSM state with the low-level core API.

Run:
    python examples/01_basic_save_load.py --model path/to/falcon-mamba-7b-Q4_K_M.gguf
"""

import argparse
import sys
from pathlib import Path
from llama_cpp import Llama
from memba import core

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model",      required=True,          help="Path to GGUF file")
    parser.add_argument("--gpu-layers", type=int, default=0,    help="GPU layers (0=CPU)")
    parser.add_argument("--state",      default="/tmp/demo.memb", help="State file path")
    args = parser.parse_args()

    # ── Load model ───────────────────────────────────────────────
    print(f"Loading model: {args.model}", flush=True)
    llama = Llama(
        model_path=args.model,
        n_ctx=4096,
        n_gpu_layers=args.gpu_layers,
        verbose=False,
    )

    # ── First inference ───────────────────────────────────────────
    prompt1 = "The capital of France is"
    print(f"\nPrompt 1: {prompt1!r}")
    out1 = llama(prompt1, max_tokens=32, echo=False)
    text1 = out1["choices"][0]["text"].strip()
    print(f"Response: {text1}")
    print(f"State size before save: {core.get_state_size(llama):,} bytes")

    # ── Save state ────────────────────────────────────────────────
    print(f"\nSaving state to {args.state} …")
    core.save_state(llama, args.model, args.state)
    saved_bytes = Path(args.state).stat().st_size
    print(f"Saved ({saved_bytes:,} bytes on disk)")

    # ── Second inference — accumulates on top of state 1 ─────────
    prompt2 = " Its population is approximately"
    print(f"\nPrompt 2 (continuous): {prompt2!r}")
    out2 = llama(prompt2, max_tokens=24, echo=False)
    text2 = out2["choices"][0]["text"].strip()
    print(f"Response: {text2}")

    # ── Reload the saved state ────────────────────────────────────
    print(f"\nRestoring state from checkpoint …")
    core.load_state(llama, args.model, args.state)
    print("State restored.")

    # ── Same prompt 2 again — should reproduce same answer ────────
    print(f"\nPrompt 2 again (after restore): {prompt2!r}")
    out3 = llama(prompt2, max_tokens=24, echo=False)
    text3 = out3["choices"][0]["text"].strip()
    print(f"Response: {text3}")

    match = text2 == text3
    print(f"\nReproducible? {'YES' if match else 'NO (expected for non-greedy sampling)'}")
    print("\nDone.")

if __name__ == "__main__":
    main()