import argparse
from time import perf_counter

from semantic_notes import api

CASES = {
    "basic": (
        "A workshop has 18 seats. Four rows of three seats are taken. "
        "How many seats are free? Put just the number in your final answer.",
        6,
    ),
    "extra-seat": (
        "A workshop has 18 seats. Four rows of three seats are taken, and one additional "
        "seat is taken. How many seats are free? Put just the number in your final answer.",
        5,
    ),
}

parser = argparse.ArgumentParser(description="Compare two thinking settings on a small math problem.")
parser.add_argument("case", choices=CASES, default="basic", nargs="?")
args = parser.parse_args()
prompt, expected = CASES[args.case]
print(f"{prompt}\nExpected answer: {expected}\n")

try:
    for think in (False, True):
        started = perf_counter()
        result = api("/api/chat", {
            "model": "qwen3:0.6b",
            "messages": [{"role": "user", "content": prompt}],
            "think": think,
            "stream": False,
            "options": {"temperature": 0, "num_ctx": 2048, "num_predict": 1024},
        })
        message = result["message"]
        elapsed = perf_counter() - started
        print(f"Thinking requested: {think}")
        print(f"  Visible answer: {message.get('content') or '(empty)'}")
        print(f"  Thinking characters: {len(message.get('thinking') or '')}")
        print(f"  Generated tokens: {result.get('eval_count', 'unreported')}")
        print(f"  Stop reason: {result.get('done_reason', 'unreported')}")
        print(f"  Client elapsed time: {elapsed:.2f} s\n")
except (OSError, KeyError, TypeError, ValueError) as error:
    raise SystemExit(f"Could not run the comparison: {error}") from None