from __future__ import annotations

import argparse
import asyncio
import time

from server import (
    END_SOLUTION,
    THINKING_TOKEN_BUDGETS,
    ChatMessage,
    ChatRequest,
    default_model_id,
    inference_config,
    stream_model_once,
)


DEFAULT_PROMPT = "Find the value of x if 3x + 7 = 31. Show the reasoning briefly."


async def benchmark(args: argparse.Namespace) -> None:
    request = ChatRequest(
        messages=[ChatMessage(role="user", content=args.prompt)],
        model=args.model or default_model_id(),
        temperature=0.0 if args.fast_mode else args.temperature,
        max_tokens=args.max_tokens,
        thinking_effort=args.thinking_effort,
        fast_mode=args.fast_mode,
    )
    generation_state: dict[str, object] = {}
    max_tokens = args.max_tokens
    if args.fast_mode:
        max_tokens += THINKING_TOKEN_BUDGETS[request.thinking_effort]

    started = time.monotonic()
    chunks: list[str] = []
    async for chunk in stream_model_once(
        request,
        max_tokens=max_tokens,
        state=generation_state,
        stop_sequences=(END_SOLUTION,),
        temperature_override=0.0 if args.fast_mode else None,
    ):
        chunks.append(chunk)
    elapsed = time.monotonic() - started

    generated_tokens = int(generation_state.get("generated_tokens") or 0)
    tokens_per_second = generated_tokens / elapsed if elapsed > 0 else 0.0
    config = inference_config()
    print(f"backend={config.backend}")
    print(f"model={request.model}")
    print(f"fast_mode={request.fast_mode}")
    print(f"generated_tokens={generated_tokens}")
    print(f"seconds={elapsed:.3f}")
    print(f"tokens_per_second={tokens_per_second:.2f}")
    print(f"stop_reason={generation_state.get('stop_reason')}")
    print(f"text_chars={sum(len(chunk) for chunk in chunks)}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Benchmark one Transformers streaming pass.")
    parser.add_argument("--model", default=None, help="Model id from /api/models. Defaults to server default.")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--max-tokens", type=int, default=512)
    parser.add_argument("--temperature", type=float, default=0.6)
    parser.add_argument("--thinking-effort", default="speed", choices=sorted(THINKING_TOKEN_BUDGETS))
    parser.add_argument("--fast-mode", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    asyncio.run(benchmark(parse_args()))