#!/usr/bin/env python3
"""Interactive Terminal Chat for `o1` (3.32M params, 32,512 B O(1) carrier + online GSSM)."""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import torch

HERE = Path(__file__).resolve().parent
if str(HERE) not in sys.path:
    sys.path.insert(0, str(HERE))

from o1_native.organism import O1Organism  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser(description="Interactive Terminal Chat with o1")
    ap.add_argument("--checkpoint", type=str, default=str(HERE / "o1.pt"))
    ap.add_argument("--device", type=str, default="cpu", help="cpu is fastest for B=1 chat; cuda/mps also supported")
    ap.add_argument("--no-learn", action="store_true", help="Disable live per-turn online weight learning")
    ap.add_argument("--max-tokens", type=int, default=96)
    ap.add_argument("--temperature", type=float, default=0.0)
    args = ap.parse_args()

    ckpt = Path(args.checkpoint)
    payload = torch.load(ckpt, map_location="cpu", weights_only=False)
    meta = payload.get("meta", {})
    org = O1Organism.load_snapshot(ckpt, device_override=args.device)
    learn_mode = not args.no_learn

    n_params = sum(p.numel() for p in org.model.parameters())
    print(f"o1 | params={n_params:,} | carrier={org.carrier.nbytes():,} B | learn={'on' if learn_mode else 'off'} | dev={args.device}")
    if meta:
        print(f"by {meta.get('author', 'David Tom Foss')} | {meta.get('email', 'david@wirku.ng')} | {meta.get('x', 'https://x.com/fossdt')}")
    print("Commands: /reset  /learn on|off  /stats  /save <path>  /exit\n")

    while True:
        try:
            msg = input("you> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not msg:
            continue
        if msg in ("/exit", "/quit"):
            break
        if msg == "/reset":
            org.reset_carrier()
            print("[o1] carrier reset to zeros (32,512 B)")
            continue
        if msg == "/learn on":
            learn_mode = True
            print("[o1] live online learning: ON")
            continue
        if msg == "/learn off":
            learn_mode = False
            print("[o1] live online learning: OFF")
            continue
        if msg == "/stats":
            wb = org.causal_store.verify_weltbuch()
            print(
                f"[o1] tokens_seen={org.total_tokens_seen:,} | plastic_updates={org.plastic_updates} "
                f"| carrier_bytes={org.carrier.nbytes():,} | 16d_drift={org.invariant_16d_drift():.2e} "
                f"| weltbuch_valid={wb.get('valid')} ({wb.get('total_triples')} triples)"
            )
            continue
        if msg.startswith("/save "):
            target = msg[6:].strip()
            if target:
                org.save_snapshot(target)
                print(f"[o1] saved state -> {target}")
            continue

        turn_prompt = f"User: {msg}\nAssistant:"
        t0 = time.perf_counter()
        out = org.generate(
            turn_prompt,
            max_new_tokens=args.max_tokens,
            temperature=args.temperature,
            use_live_carrier=True,
            stop_on_newline=False,
        )
        ans = str(out.get("answer", "")).split("\nUser:", 1)[0].strip()
        dt = max(time.perf_counter() - t0, 1e-6)
        toks = max(int(out.get("emitted_token_count", 0)), 1)

        step_rep = org.stream_step(turn_prompt, learn=learn_mode)
        _ = org.stream_step(f" {ans}\n", learn=False)
        route = out.get("organ_route") or "TEXT"
        loss_nats = float(step_rep.get("metrics", {}).get("loss_nats", 0.0))
        print(f"o1 [{route} | {toks/dt:.1f} tok/s | loss={loss_nats:.3f}]> {ans}\n")


if __name__ == "__main__":
    main()