from __future__ import annotations

import math
from statistics import mean

from .models import Request, RequestMetrics, SimulationConfig


def percentile(values: list[float], q: float) -> float:
    if not values:
        return 0.0
    vals = sorted(values)
    if len(vals) == 1:
        return vals[0]
    pos = (len(vals) - 1) * q
    lo = math.floor(pos)
    hi = math.ceil(pos)
    if lo == hi:
        return vals[lo]
    frac = pos - lo
    return vals[lo] * (1.0 - frac) + vals[hi] * frac


def request_metrics(req: Request, cfg: SimulationConfig) -> RequestMetrics:
    if req.first_token_time is None or req.completion_time is None:
        raise ValueError("Request is incomplete")
    ttft_ms = (req.first_token_time - req.arrival_time) * 1000.0
    e2e_ms = (req.completion_time - req.arrival_time) * 1000.0
    if req.output_tokens <= 1:
        tpot_ms = 0.0
    else:
        tpot_ms = (req.completion_time - req.first_token_time) * 1000.0 / (req.output_tokens - 1)
    prefill_start = req.first_prefill_time if req.first_prefill_time is not None else req.arrival_time
    queue_ms = max(0.0, (prefill_start - req.arrival_time) * 1000.0)
    met_ttft = ttft_ms <= cfg.slo_ttft_ms
    met_e2e = e2e_ms <= cfg.slo_e2e_ms
    return RequestMetrics(
        request_id=req.request_id,
        arrival_time=req.arrival_time,
        prompt_tokens=req.prompt_tokens,
        output_tokens=req.output_tokens,
        ttft_ms=ttft_ms,
        e2e_ms=e2e_ms,
        tpot_ms=tpot_ms,
        queue_ms=queue_ms,
        met_ttft_slo=met_ttft,
        met_e2e_slo=met_e2e,
        met_all_slos=met_ttft and met_e2e,
    )


def summarize(completed: list[Request], cfg: SimulationConfig, makespan_s: float, busy_time_s: float) -> tuple[dict, dict]:
    metrics = [request_metrics(r, cfg) for r in completed]
    ttft = [m.ttft_ms for m in metrics]
    e2e = [m.e2e_ms for m in metrics]
    tpot = [m.tpot_ms for m in metrics]
    queue = [m.queue_ms for m in metrics]
    total_output = sum(r.output_tokens for r in completed)
    met = sum(m.met_all_slos for m in metrics)
    met_ttft = sum(m.met_ttft_slo for m in metrics)
    met_e2e = sum(m.met_e2e_slo for m in metrics)
    ttft_only_fail = sum((not m.met_ttft_slo) and m.met_e2e_slo for m in metrics)
    e2e_only_fail = sum(m.met_ttft_slo and (not m.met_e2e_slo) for m in metrics)
    both_fail = sum((not m.met_ttft_slo) and (not m.met_e2e_slo) for m in metrics)
    duration = max(makespan_s, 1e-9)
    count = len(metrics)

    summary = {
        "requests_completed": len(completed),
        "request_throughput_rps": len(completed) / duration,
        "output_throughput_tps": total_output / duration,
        "goodput_rps": met / duration,
        "slo_attainment": met / count if count else 0.0,
        "ttft_slo_attainment": met_ttft / count if count else 0.0,
        "e2e_slo_attainment": met_e2e / count if count else 0.0,
        "slo_failure_breakdown": {"ttft_only": ttft_only_fail, "e2e_only": e2e_only_fail, "both": both_fail},
        "simulated_makespan_s": makespan_s,
        "busy_fraction": min(1.0, busy_time_s / duration),
        "mean_prompt_tokens": mean([r.prompt_tokens for r in completed]) if completed else 0.0,
        "mean_output_tokens": mean([r.output_tokens for r in completed]) if completed else 0.0,
    }
    latency = {
        "ttft_ms": {"p50": percentile(ttft, 0.50), "p95": percentile(ttft, 0.95), "p99": percentile(ttft, 0.99)},
        "e2e_ms": {"p50": percentile(e2e, 0.50), "p95": percentile(e2e, 0.95), "p99": percentile(e2e, 0.99)},
        "tpot_ms": {"p50": percentile(tpot, 0.50), "p95": percentile(tpot, 0.95), "p99": percentile(tpot, 0.99)},
        "queue_ms": {"p50": percentile(queue, 0.50), "p95": percentile(queue, 0.95), "p99": percentile(queue, 0.99)},
    }
    return summary, latency