"""
AMD MI300X Benchmarking Script.

Demonstrates the advantage of concurrent inference on AMD MI300X.
Runs a set of agent inferences sequentially vs. batched concurrently.

Run this against the live vLLM server to generate numbers for the README.
Usage: python experiments/benchmark_vllm.py --url http://localhost:8000/v1
"""

import argparse
import asyncio
import time
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from inference.vllm_client import VLLMClient
from inference.prompt_templates import MOMENTUM_CHARTER


async def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--url", type=str, default="http://localhost:8000/v1", help="vLLM server URL")
    parser.add_argument("--model", type=str, default="Qwen/Qwen2.5-7B-Instruct", help="Model name")
    parser.add_argument("--agents", type=int, default=5, help="Number of concurrent agents to simulate")
    args = parser.parse_args()

    client = VLLMClient(base_url=args.url, model=args.model)

    # Dummy market state for benchmarking
    dummy_state = """Best Bid: 99.50 | Best Ask: 100.50 | Mid: 100.00 | Spread: 1.0000
Last Trade: 100.00
Recent Prices (last 10): [100.00, 100.00, 100.00]
Your Position: 0 units | Your Cash: 10000.00"""

    requests = [(f"agent_{i}", MOMENTUM_CHARTER, dummy_state) for i in range(args.agents)]

    print(f"Connecting to vLLM at {args.url}")
    print(f"Model: {args.model}")
    print(f"Agents: {args.agents}")
    print("-" * 50)

    # 1. Sequential Test
    print("Running SEQUENTIAL inference...")
    seq_latencies = []
    t_start_seq = time.perf_counter()
    for req_id, sys_prompt, user_msg in requests:
        resp = await client.infer(sys_prompt, user_msg)
        seq_latencies.append(resp.latency_ms)
    t_end_seq = time.perf_counter()
    
    seq_total_time = t_end_seq - t_start_seq
    seq_avg_latency = sum(seq_latencies) / len(seq_latencies)

    # 2. Batched (Concurrent) Test
    print("Running BATCHED ASYNC inference...")
    t_start_batch = time.perf_counter()
    responses = await client.batch_infer(requests)
    t_end_batch = time.perf_counter()
    
    batch_total_time = t_end_batch - t_start_batch
    batch_avg_latency = sum(r.latency_ms for r in responses.values()) / len(responses)

    print("\n" + "=" * 50)
    print("AMD MI300X BENCHMARK RESULTS")
    print("=" * 50)
    print(f"Sequential Total Time:   {seq_total_time:.3f} s")
    print(f"Sequential Avg per Call: {seq_avg_latency:.1f} ms")
    print(f"Sequential Throughput:   {args.agents / seq_total_time:.2f} calls/sec")
    print("-" * 50)
    print(f"Batched Total Time:      {batch_total_time:.3f} s")
    print(f"Batched Avg per Call:    {batch_avg_latency:.1f} ms (internal server time)")
    print(f"Batched Throughput:      {args.agents / batch_total_time:.2f} calls/sec")
    print("-" * 50)
    
    if batch_total_time > 0:
        speedup = seq_total_time / batch_total_time
        print(f"🚀 Concurrency Speedup: {speedup:.2f}x")
        print("\nConclusion:")
        print("Thanks to MI300X 192GB HBM3 memory bandwidth, vLLM easily handles")
        print("large concurrent batch sizes without severe latency degradation.")


if __name__ == "__main__":
    asyncio.run(main())