#!/usr/bin/env python3
"""Terminal demo for CBT-18.95M. Sums like `123+456` go to the adder; anything else is a chat turn.

    python chat.py                          # interactive
    python chat.py "98765432109876+1234"    # one prompt
    python chat.py --online "Who was Ada Lovelace?"
    python chat.py --variant twin_1d "31415926535897932384626433832795+27182818284590452353602874713527"
"""

import argparse
import functools
import re

from cbt import add, chat, execute_tool, load

SUM = re.compile(r"^\s*(\d+)\s*\+\s*(\d+)\s*=?\s*$")


def respond(model, tok, text: str, online: bool, show_raw: bool) -> None:
    m = SUM.match(text)
    if m:
        r = add(model, tok, m.group(1), m.group(2))
        print(f"  scratchpad : {r['scratchpad']}")
        print(f"  answer     : {r['answer']}   [{'PASS' if r['correct'] else 'FAIL'}]")
    else:
        r = chat(model, tok, text, tool_executor=functools.partial(execute_tool, online=online))
        for key in ("think", "call", "result", "answer"):
            if r[key] is not None:
                print(f"  {key:<10} : {r[key]}")
    if show_raw:
        print(f"  raw        : {r['raw']}")
    print()


def main() -> None:
    p = argparse.ArgumentParser(description="CBT-18.95M terminal demo")
    p.add_argument("prompt", nargs="?", help="single prompt; omit for interactive mode")
    p.add_argument("--model", default=None, help="local folder or Hugging Face repo id (default: this folder)")
    p.add_argument("--variant", default="cbt_2d", choices=["cbt_2d", "twin_1d"], help="released 2D model or its 1D-RoPE twin")
    p.add_argument("--device", default=None, help="cpu or cuda (default: auto)")
    p.add_argument("--online", action="store_true", help="let the search tool query Wikipedia")
    p.add_argument("--raw", action="store_true", help="also print the raw token stream")
    args = p.parse_args()

    model, tok = load(args.model, args.device, args.variant)
    if args.prompt:
        respond(model, tok, args.prompt, args.online, args.raw)
        return

    print(f"CBT-18.95M ({args.variant}). Type a sum (e.g. 123456789+987654321) or a question. Ctrl-D to quit.\n")
    while True:
        try:
            text = input("> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if text.lower() in ("exit", "quit"):
            break
        if text:
            respond(model, tok, text, args.online, args.raw)


if __name__ == "__main__":
    main()