"""
02_chat_session.py — multi-turn chat that persists across Python processes.

First run  : model has no prior context.
Second run : model continues from the saved SSM state.

Run twice:
    python examples/02_chat_session.py --model path/to/falcon-mamba-7b-Q4_K_M.gguf
    python examples/02_chat_session.py --model path/to/falcon-mamba-7b-Q4_K_M.gguf
"""

import argparse
from memba import Session

TURNS = [
    "My name is Alex. I am a researcher studying ancient Roman aqueducts.",
    "What is the most famous aqueduct I should know about?",
    "How long did it take to build?",
]

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("--session",    default="aqueduct_research")
    parser.add_argument("--state-dir",  default="~/.memba/states")
    args = parser.parse_args()

    print(f"Session: {args.session!r}")
    print(f"Model  : {args.model}")
    print("-" * 60)

    sess = Session(
        model_path=args.model,
        session_id=args.session,
        state_dir=args.state_dir,
        n_gpu_layers=args.gpu_layers,
        verbose=False,
    )
    print(f"State size on load: {sess.state_size:,} bytes\n")

    # Only feed the turns that haven't been answered yet.
    # A real app would track which turns were already fed; here we keep it simple
    # and feed all TURNS every run — the SSM state update is idempotent in terms
    # of demonstrating cross-process continuity.
    for i, turn in enumerate(TURNS, 1):
        print(f"[Turn {i}] User: {turn}")
        reply = sess.chat(turn, max_tokens=200)
        print(f"[Turn {i}] Model: {reply}")
        print()

    saved_path = sess.save()
    print(f"State saved → {saved_path}")
    print(f"State size  : {sess.state_size:,} bytes")
    print("\nRun this script again — the model will continue from this checkpoint.")

if __name__ == "__main__":
    main()