"""ChatEngine checks: prompt lookup decoding == plain greedy; prefix reuse == recompute. Prints speeds.
usage: python test_chat.py ENGINE_DIR [DEVICE]"""
import sys, time
from onw.chat import ChatEngine

CODE = '''def load_config(path):
    with open(path) as f:
        data = json.load(f)
    if "name" not in data:
        raise ValueError("missing name")
    data["port"] = int(data.get("port", 8080))
    data["debug"] = bool(data.get("debug", False))
    return data'''
TEXT = ("東京は日本の首都であり、政治・経済・文化の中心地である。江戸時代には徳川幕府が置かれ、1868年の明治維新で東京と"
        "改称された。現在は約1400万人が暮らす世界有数の大都市で、多くの企業や大学が集まっている。")
PROMPTS = {
    "code edit": f"次のPython関数に型ヒントを付けて、関数全体を書き直してください。説明は不要です。\n\n```python\n{CODE}\n```",
    "summary": f"次の文章を、元の表現をなるべく使って2文に要約してください。\n\n{TEXT}",
    "free chat": "NPUとGPUの違いを、身近なたとえを使って説明してください。",
}


def run(e, msgs, n=160):
    parts, st = [], None
    for d in e.stream_chat([dict(m) for m in msgs], n):
        if isinstance(d, dict):
            st = d
        else:
            parts.append(d)
    return "".join(parts), st


def main():
    e = ChatEngine(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else None)
    for name, q in PROMPTS.items():
        msgs = [{"role": "user", "content": q}]
        auto = e.pld
        e.pld, e.checkpoint = False, None
        a, s1 = run(e, msgs)
        if not auto:
            print(f"{name:10s} plain {s1['decode_tok_s']:5.1f} tok/s ({s1['completion_tokens']} tokens) | prefill "
                  f"{s1['prompt_tokens']} tok {s1['prefill_ms']:.0f} ms | PLD off for this model", flush=True)
            continue
        e.pld, e.checkpoint = True, None
        b, s2 = run(e, msgs)
        p = s2["pld"]
        print(f"{name:10s} plain {s1['decode_tok_s']:5.1f} tok/s | PLD {s2['decode_tok_s']:5.1f} tok/s "
              f"({p['blocks']} blocks, {p['accepted']} accepted of {s2['completion_tokens']}) | identical={a == b} "
              f"| prefill {s1['prompt_tokens']} tok {s1['prefill_ms']:.0f} ms", flush=True)
    t1 = [{"role": "user", "content": "日本の首都はどこですか？一文で。"}]
    e.checkpoint = None
    a1, _ = run(e, t1, 60)
    t2 = t1 + [{"role": "assistant", "content": a1}, {"role": "user", "content": "その都市の人口は？一文で。"}]
    a2, s2 = run(e, t2, 60)
    e.checkpoint = None
    a3, s3 = run(e, t2, 60)
    print(f"multi-turn: reuse {s2['cached_tokens']}/{s2['prompt_tokens']} tokens, prefill {s2['prefill_ms']:.0f} ms "
          f"vs {s3['prefill_ms']:.0f} ms from scratch; identical={a2 == a3}\n  turn 1: {a1!r}\n  turn 2: {a2!r}")


if __name__ == "__main__":
    main()