"""5.6 推理的两段与缓存（本地版）。对一个本地 llama-server 发请求，读它返回的 timings：
prompt_n 是这次真正重新计算的输入 token，cache_n 是直接复用缓存的 token，
prompt_ms 是 prefill 用的时间，predicted_ms / predicted_n 是 decode 每个 token 的时间。

先在另一个终端起服务（第一次会自动下载约 1.2 GB 的模型）：
    llama-server -hf unsloth/Qwen3.5-2B-GGUF:Q4_K_M --no-mmproj -c 65536 -np 1 --port 8080
再跑：
    python local_cache.py
只用标准库，不需要装任何包。
"""
import argparse
import json
import random
import string
import time
import urllib.request
from datetime import datetime
from pathlib import Path

URL = 'http://127.0.0.1:8080'


def complete(prompt, n_predict):
    body = json.dumps({'prompt': prompt, 'n_predict': n_predict, 'cache_prompt': True, 'temperature': 0}).encode()
    req = urllib.request.Request(URL + '/completion', data=body, headers={'Content-Type': 'application/json'})
    r = json.load(urllib.request.urlopen(req, timeout=600))
    return r['content'], r['timings']


def filler(n_lines, seed):
    rnd = random.Random(seed)
    word = lambda: ''.join(rnd.choice(string.ascii_lowercase) for _ in range(rnd.randint(3, 8)))
    return '\n'.join(f'{i}: ' + ' '.join(word() for _ in range(12)) for i in range(n_lines))


def exp_input():
    print('\n## 实验 1：输入变长，只生成 8 个 token\n')
    print('| 输入 token | prefill 用时 | prefill 速度 |')
    print('|---|---|---|')
    out = []
    for n_lines in [25, 100, 400, 1200]:
        _, t = complete(filler(n_lines, f'in-{n_lines}-{time.time()}') + '\nHow many lines are above? Answer:', 8)
        out.append(t)
        print(f'| {t["prompt_n"]:,} | {t["prompt_ms"] / 1000:.2f} 秒 | {t["prompt_per_second"]:,.0f} token/秒 |')
    return out


def exp_output():
    print('\n## 实验 2：输入很短，输出变长\n')
    print('| 输出 token | decode 用时 | 每个输出 token | decode 速度 |')
    print('|---|---|---|---|')
    out = []
    for n in [50, 200, 800]:
        _, t = complete(f'[{time.time()}] Count from 1 upward, separated by commas: 1, 2, 3,', n)
        out.append(t)
        print(f'| {t["predicted_n"]} | {t["predicted_ms"] / 1000:.2f} 秒 | {t["predicted_per_token_ms"]:.1f} 毫秒 '
              f'| {t["predicted_per_second"]:.0f} token/秒 |')
    return out


QUESTIONS = ['What is the first word on line 3?', 'How many words are on line 10?',
             'What is the last word on line 25?', 'What is the second word on line 7?',
             'What is the first word on line 40?', 'What is the third word on line 12?']


def exp_cache(rounds, doc_text):
    print(f'\n## 实验 3：同一份长文档上连问 {rounds} 轮\n')
    res = {}
    for mode in ['append_only', 'timestamp_first']:
        doc = doc_text or filler(1200, f'{mode}-{time.time()}')
        convo, rows = '', []
        for i in range(rounds):
            head = '' if mode == 'append_only' else f'Current time: {datetime.now().isoformat()}\n\n'
            convo += f'\nQ: {QUESTIONS[i % len(QUESTIONS)]}\nA:'
            reply, t = complete(head + doc + convo, 16)
            convo += reply
            rows.append(t)
        res[mode] = rows
    print('| 轮次 | 追加：重算 / 复用 token | 追加：prefill 用时 | 时间戳：重算 / 复用 token | 时间戳：prefill 用时 |')
    print('|---|---|---|---|---|')
    for i in range(rounds):
        a, b = res['append_only'][i], res['timestamp_first'][i]
        print(f'| {i + 1} | {a["prompt_n"]:,} / {a["cache_n"]:,} | {a["prompt_ms"] / 1000:.2f} 秒 '
              f'| {b["prompt_n"]:,} / {b["cache_n"]:,} | {b["prompt_ms"] / 1000:.2f} 秒 |')
    tot = {m: sum(t['prompt_n'] for t in res[m]) for m in res}
    print(f'\n{rounds} 轮一共重新计算的输入 token：只在末尾追加 {tot["append_only"]:,}，开头插时间戳 {tot["timestamp_first"]:,}。'
          '\n用按量计费的 API 时，重新计算的这部分按未命中价收费，复用的部分按命中价（通常便宜一个数量级）。')
    return res


def main():
    global URL
    ap = argparse.ArgumentParser()
    ap.add_argument('--url', default=URL)
    ap.add_argument('--only', choices=['input', 'output', 'cache'])
    ap.add_argument('--rounds', type=int, default=6)
    ap.add_argument('--doc', help='缓存实验用你自己的长文档（文本文件）')
    ap.add_argument('--out', default='results')
    args = ap.parse_args()
    URL = args.url.rstrip('/')
    props = json.load(urllib.request.urlopen(URL + '/props'))
    model = Path(props.get('model_path', '?')).name
    print(f'本地模型 {model}，上下文 {props.get("default_generation_settings", {}).get("n_ctx")}，{datetime.now():%Y-%m-%d %H:%M}')
    doc_text = Path(args.doc).read_text(encoding='utf-8') if args.doc else None
    res = {'model': model}
    if args.only in (None, 'input'):
        res['input'] = exp_input()
    if args.only in (None, 'output'):
        res['output'] = exp_output()
    if args.only in (None, 'cache'):
        res['cache'] = exp_cache(args.rounds, doc_text)
    Path(args.out).mkdir(exist_ok=True)
    path = Path(args.out) / f'local_{datetime.now():%Y%m%d_%H%M}.json'
    path.write_text(json.dumps(res, ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'\n原始 timings 已存到 {path}')


if __name__ == '__main__':
    main()