"""5.6 The two phases of inference and caching (API version). Three streaming timing experiments:

1. Varying input length: fix output to a few tokens and measure how Time to First Token (TTFT) grows with input length (dominated by prefill).
2. Varying output length: keep input short and measure how total time grows with output length and how long each output token takes (decode).
3. Prefix caching: ask multiple turns of questions over the same long document, comparing append-only against prepending a current timestamp
   on every turn by measuring Time to First Token on each turn.

Usage:
    export AI_BUILDER_TOKEN=sk_...
    python api_timing.py                     # Default: deepseek, runs all 3 experiments (~5 min)
    python api_timing.py --model grok-4-fast
    python api_timing.py --only cache        # Run just one: input / output / cache
    python api_timing.py --doc my_notes.txt  # Use your own long document for the cache experiment
"""
import argparse
import json
import os
import random
import statistics
import string
import time
from datetime import datetime
from pathlib import Path

from openai import OpenAI

BASE_URL = os.getenv('AI_BUILDER_SPACE_BASE_URL', 'https://space.ai-builders.com/backend/v1')
PRICES = {'deepseek': (0.15, 0.60), 'grok-4-fast': (1.25, 2.50)}   # USD per 1M tokens (input, output), Oct 2026 AI Builder Space pricing

client = None


def stream(model, messages, max_tokens):
    """Send one streaming request and return (TTFT in seconds, total seconds, chunks received, text). Chunks roughly equal output tokens."""
    t0 = time.time()
    first, chunks, text = None, 0, []
    s = client.chat.completions.create(model=model, messages=messages, max_tokens=max_tokens, stream=True)
    for ch in s:
        if ch.choices and ch.choices[0].delta and ch.choices[0].delta.content:
            if first is None:
                first = time.time() - t0
            chunks += 1
            text.append(ch.choices[0].delta.content)
    return first, time.time() - t0, chunks, ''.join(text)


def prompt_tokens(model, messages):
    """Make one non-streaming call solely to read prompt_tokens from usage (streaming responses omit usage)."""
    r = client.chat.completions.create(model=model, messages=messages, max_tokens=1)
    return r.usage.prompt_tokens


def filler(n_words, seed):
    """Generate random English-like "words". A fresh seed per call ensures no two requests share a prefix or hit the cache."""
    rnd = random.Random(seed)
    words = [''.join(rnd.choice(string.ascii_lowercase) for _ in range(rnd.randint(3, 8))) for _ in range(n_words)]
    lines = [' '.join(words[i:i + 12]) for i in range(0, n_words, 12)]
    return '\n'.join(f'{i}: {l}' for i, l in enumerate(lines))


def exp_input(model, reps):
    print('\n## Experiment 1: Varying input length, output capped at 8 tokens\n')
    print('| Input tokens | Median TTFT | Individual runs |')
    print('|---|---|---|')
    out = []
    for n_words in [500, 2500, 10000, 40000]:
        ttfts = []
        for k in range(reps):
            seed = f'{time.time()}-{k}'
            msgs = [{'role': 'user', 'content': filler(n_words, seed) + '\n\n上面一共有几行？只回答数字。'}]
            ttfts.append(stream(model, msgs, 8)[0])
        n_tok = prompt_tokens(model, [{'role': 'user', 'content': filler(n_words, 'count') + '\n\n上面一共有几行？只回答数字。'}])
        out.append(dict(prompt_tokens=n_tok, ttft=ttfts))
        print(f'| {n_tok:,} | {statistics.median(ttfts):.2f} s | {", ".join(f"{x:.2f}" for x in ttfts)} |')
    return out


def exp_output(model, reps):
    print('\n## Experiment 2: Short fixed input, varying output length\n')
    print('| Requested output tokens | Chunks received | TTFT | Total time | Time per output token |')
    print('|---|---|---|---|---|')
    out = []
    msgs = [{'role': 'user', 'content': '从 1 开始一直往上数，用逗号分隔，不要停，不要说别的。'}]
    for max_tokens in [50, 200, 800]:
        runs = [stream(model, msgs, max_tokens) for _ in range(reps)]
        ttft = statistics.median(r[0] for r in runs)
        total = statistics.median(r[1] for r in runs)
        chunks = statistics.median(r[2] for r in runs)
        per_tok = statistics.median((r[1] - r[0]) / max(r[2] - 1, 1) for r in runs)
        out.append(dict(max_tokens=max_tokens, ttft=ttft, total=total, chunks=chunks, sec_per_token=per_tok))
        print(f'| {max_tokens} | {chunks:.0f} | {ttft:.2f} s | {total:.2f} s | {per_tok * 1000:.1f} ms |')
    return out


QUESTIONS = ['第 3 行的第一个词是什么？只回答这个词。', '第 10 行一共有几个词？只回答数字。',
             '第 25 行的最后一个词是什么？只回答这个词。', '第 7 行的第二个词是什么？只回答这个词。',
             '第 40 行的第一个词是什么？只回答这个词。', '第 12 行的第三个词是什么？只回答这个词。']


def exp_cache(model, rounds, doc_text):
    print(f'\n## Experiment 3: Asking {rounds} turns over the same long document\n')
    results = {}
    for mode in ['append_only', 'timestamp_first']:
        doc = doc_text or filler(30000, f'{mode}-{time.time()}')       # Use a separate document per mode so one cannot hit the other's cache
        history, ttfts = [], []
        for i in range(rounds):
            system = doc if mode == 'append_only' else f'当前时间：{datetime.now().isoformat()}\n\n{doc}'
            q = {'role': 'user', 'content': QUESTIONS[i % len(QUESTIONS)]}
            ttft, _, _, reply = stream(model, [{'role': 'system', 'content': system}] + history + [q], 20)
            history += [q, {'role': 'assistant', 'content': reply}]
            ttfts.append(ttft)
        results[mode] = ttfts
    n_tok = prompt_tokens(model, [{'role': 'system', 'content': doc_text or filler(30000, 'count')},
                                  {'role': 'user', 'content': QUESTIONS[0]}])
    print(f'Input per turn: ~{n_tok:,} tokens.\n')
    print('| Turn | Append-only | Prepend timestamp |')
    print('|---|---|---|')
    for i in range(rounds):
        print(f'| {i + 1} | {results["append_only"][i]:.2f} s | {results["timestamp_first"][i]:.2f} s |')
    a, b = results['append_only'][1:], results['timestamp_first'][1:]
    print(f'| Median from Turn 2 onward | {statistics.median(a):.2f} s | {statistics.median(b):.2f} s |')
    pin, pout = PRICES.get(model, (0, 0))
    print(f'\nAt AI Builder Space rates, input costs ~${n_tok * pin / 1e6:.4f} per turn under both patterns: '
          'the gateway does not discount cache hits or report cached token counts in usage. Caching here saves time.')
    return dict(prompt_tokens=n_tok, **results)


def main():
    global client
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='deepseek', help='Must support streaming: deepseek and grok-4-fast work; kimi and gemini do not')
    ap.add_argument('--only', choices=['input', 'output', 'cache'])
    ap.add_argument('--reps', type=int, default=3)
    ap.add_argument('--rounds', type=int, default=6)
    ap.add_argument('--doc', help='Text file to use as the long document in the cache experiment; defaults to ~100K tokens of random text')
    ap.add_argument('--out', default='results')
    args = ap.parse_args()
    token = os.getenv('AI_BUILDER_TOKEN')
    if not token:
        raise SystemExit('Please set the AI_BUILDER_TOKEN environment variable (your AI Builder Space key)')
    client = OpenAI(base_url=BASE_URL, api_key=token, timeout=600)
    doc_text = Path(args.doc).read_text(encoding='utf-8') if args.doc else None

    print(f'Model {args.model}, {datetime.now():%Y-%m-%d %H:%M}')
    res = {'model': args.model}
    if args.only in (None, 'input'):
        res['input'] = exp_input(args.model, args.reps)
    if args.only in (None, 'output'):
        res['output'] = exp_output(args.model, args.reps)
    if args.only in (None, 'cache'):
        res['cache'] = exp_cache(args.model, args.rounds, doc_text)
    Path(args.out).mkdir(exist_ok=True)
    path = Path(args.out) / f'api_{args.model}_{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()