"""
Naive Batching: process one request at a time, no concurrency.
This is the baseline — every LLM serving system starts here.
"""

import time
import torch
import numpy as np
from dataclasses import dataclass
from typing import List
from transformers import AutoTokenizer, AutoModelForCausalLM


@dataclass
class InferenceResult:
    prompt: str
    output: str
    input_tokens: int
    output_tokens: int
    latency_ms: float
    tokens_per_second: float


class NaiveBatchingEngine:
    """
    Sequential inference: each request waits for the previous to complete.
    Problems:
    - GPU sits idle between requests
    - No sharing of KV cache computation
    - Latency scales linearly with queue depth
    """

    def __init__(self, model_name: str, device: str = "auto"):
        print(f"[NaiveBatching] Loading {model_name}...")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map=device,
        )
        self.model.eval()
        self.device = next(self.model.parameters()).device
        print(f"[NaiveBatching] Model loaded on {self.device}")

    @torch.no_grad()
    def generate_single(self, prompt: str, max_new_tokens: int = 50) -> InferenceResult:
        """Generate for a single prompt, sequentially."""
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
        input_len = inputs["input_ids"].shape[1]

        start = time.perf_counter()
        output_ids = self.model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=self.tokenizer.eos_token_id,
        )
        elapsed_ms = (time.perf_counter() - start) * 1000

        output_len = output_ids.shape[1] - input_len
        output_text = self.tokenizer.decode(
            output_ids[0][input_len:], skip_special_tokens=True
        )
        tps = (output_len / elapsed_ms) * 1000

        return InferenceResult(
            prompt=prompt,
            output=output_text,
            input_tokens=input_len,
            output_tokens=output_len,
            latency_ms=elapsed_ms,
            tokens_per_second=tps,
        )

    def benchmark(
        self, prompts: List[str], max_new_tokens: int = 50
    ) -> dict:
        """Run prompts sequentially and collect latency statistics."""
        results = []
        for i, prompt in enumerate(prompts):
            result = self.generate_single(prompt, max_new_tokens)
            results.append(result)
            print(f"  [{i+1}/{len(prompts)}] {result.latency_ms:.1f}ms, "
                  f"{result.tokens_per_second:.1f} tok/s")

        latencies = [r.latency_ms for r in results]
        tps_values = [r.tokens_per_second for r in results]
        total_time = sum(latencies)

        return {
            "method": "naive_sequential",
            "n_requests": len(prompts),
            "total_time_ms": total_time,
            "throughput_requests_per_sec": len(prompts) / (total_time / 1000),
            "throughput_tokens_per_sec": sum(r.output_tokens for r in results) / (total_time / 1000),
            "latency_p50_ms": float(np.percentile(latencies, 50)),
            "latency_p95_ms": float(np.percentile(latencies, 95)),
            "latency_p99_ms": float(np.percentile(latencies, 99)),
            "latency_mean_ms": float(np.mean(latencies)),
            "tokens_per_second_mean": float(np.mean(tps_values)),
            "results": results,
        }