"""5.6 The two phases of inference and caching (local version). Sends requests to a local llama-server and reads its timings:
prompt_n is the number of input tokens actually recomputed, cache_n is the number of tokens reused directly from cache,
prompt_ms is prefill time, and predicted_ms / predicted_n is decode time per output token.

Start the server in another terminal first (downloads the ~1.2 GB model automatically on first run):
    llama-server -hf unsloth/Qwen3.5-2B-GGUF:Q4_K_M --no-mmproj -c 65536 -np 1 --port 8080
Then run:
    python local_cache.py
Uses only the Python standard library; no extra packages required.
"""
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## Experiment 1: Varying input length, generating only 8 tokens\n')
    print('| Input tokens | Prefill time | Prefill speed |')
    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} s | {t["prompt_per_second"]:,.0f} tokens/s |')
    return out


def exp_output():
    print('\n## Experiment 2: Short fixed input, varying output length\n')
    print('| Output tokens | Decode time | Time per output token | Decode speed |')
    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} s | {t["predicted_per_token_ms"]:.1f} ms '
              f'| {t["predicted_per_second"]:.0f} tokens/s |')
    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## Experiment 3: Asking {rounds} turns over the same long document\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('| Turn | Append-only: recomputed / cached tokens | Append-only: prefill time | Timestamp: recomputed / cached tokens | Timestamp: prefill time |')
    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} s '
              f'| {b["prompt_n"]:,} / {b["cache_n"]:,} | {b["prompt_ms"] / 1000:.2f} s |')
    tot = {m: sum(t['prompt_n'] for t in res[m]) for m in res}
    print(f'\nTotal recomputed input tokens across {rounds} turns: append-only {tot["append_only"]:,}, prepend-timestamp {tot["timestamp_first"]:,}.'
          '\nOn usage-billed APIs, recomputed tokens pay the cache-miss rate while reused tokens pay the cache-hit rate (typically an order of magnitude cheaper).')
    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='Text file to use as your own long document in the cache experiment')
    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'Local model {model}, context {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'\nRaw timings saved to {path}')


if __name__ == '__main__':
    main()