"""Chat with Ling-mini-2.0 in a terminal, on the megakernel.

    python chat.py [--max-new 1024]

Greedy decoding, thinking off. Each answer re-reads only the new part of the conversation (the rest
is still in the kernel's cache). Windows lowers an idle GPU's clocks, so the kernel warms up at start
and before each answer. The speed shown is the kernel's own while it writes. An empty line or
"exit" quits.
"""

from __future__ import annotations

import argparse
import sys
import time

from ling_kernel import HERE, LingKernel

CHUNK = 32          # tokens per launch, printed as they arrive


def calibrate(mk: LingKernel, tok, budget: float = 20.0) -> float:
    ids = tok("Hello")["input_ids"]
    first = mk.forced(ids)[-1]
    start, best, since = time.perf_counter(), float("inf"), time.perf_counter()
    while time.perf_counter() - start < budget:
        s = mk._launch([first], len(ids), CHUNK, False)
        if s < 0.97 * best:
            best, since = s, time.perf_counter()
        elif time.perf_counter() - since > 1.0:
            break
    mk.cached = []
    return best


def warm_up(mk: LingKernel, first: int, pos: int, budget: float = 3.0) -> None:
    start = time.perf_counter()
    while time.perf_counter() - start < budget:
        if mk._launch([first], pos, CHUNK, False) < 1.15 * mk.best_launch:
            break


def answer(mk: LingKernel, tok, eos: set[int], messages: list[dict], max_new: int) -> str:
    ids = tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
    if hasattr(ids, "keys"):
        ids = ids["input_ids"]
    ids = list(ids)
    if len(ids) + max_new + CHUNK > mk.max_pos:
        print(f"[the conversation is too long for the {mk.max_pos}-token cache; type 'exit' and start again]")
        return ""
    n = 0
    while n < min(len(ids) - 1, len(mk.cached)) and ids[n] == mk.cached[n]:
        n += 1
    preds = mk.forced(ids[n:], pos0=n)
    warm_up(mk, preds[-1], len(ids))
    out, seconds, steps, pos, shown = [preds[-1]], 0.0, 0, len(ids), ""
    print("Ling> ", end="", flush=True)
    while out[-1] not in eos and len(out) < max_new:
        k = min(CHUNK, max_new - len(out))
        seconds += mk._launch([out[-1]], pos, k, False)
        steps += k
        pos += k
        for t in mk.state.out[:k].tolist():
            out.append(t)
            if t in eos:
                break
        text = tok.decode([t for t in out if t not in eos], skip_special_tokens=True)
        print(text[len(shown):], end="", flush=True)
        shown = text
    mk.cached = ids + out[:-1]
    speed = f"{steps / seconds:.0f} tokens/s" if seconds else "-"
    print(f"\n[{len(out)} tokens; the kernel writes {speed}; read {len(ids) - n} new prompt tokens]\n", flush=True)
    return shown.strip()


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--max-new", type=int, default=1024)
    args = ap.parse_args()
    sys.stdout.reconfigure(encoding="utf-8")
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(HERE)
    print("loading the Ling-mini-2.0 megakernel (4-bit) ...", flush=True)
    mk = LingKernel(max_pos=4096, log=lambda s: print(s, flush=True))
    eos = {i for i in (tok.convert_tokens_to_ids("<|role_end|>"), tok.eos_token_id, mk.config.get("eos_token_id"))
           if isinstance(i, int) and i >= 0}
    print("warming up the GPU ...", flush=True)
    mk.best_launch = calibrate(mk, tok)
    print(f"  {CHUNK / mk.best_launch:.0f} tokens/s at full speed", flush=True)
    history: list[dict] = []
    print("type a message; an empty line or 'exit' quits\n", flush=True)
    while True:
        try:
            text = input("you> ").strip()
        except EOFError:
            break
        if not text or text.lower() in ("exit", "quit"):
            break
        messages = history + [{"role": "user", "content": text}]
        reply = answer(mk, tok, eos, messages, args.max_new)
        history = messages + [{"role": "assistant", "content": reply}]


if __name__ == "__main__":
    main()