from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from .spec import SloSpec, ThresholdRule


@dataclass(frozen=True)
class RequestOutcome:
    request_id: str
    success: bool
    ttft_ms: float | None
    tpot_ms: float | None
    prompt_tokens: int | None
    completion_tokens: int | None
    error: str = ""
    completion_tokens_source: str = ""
    completed_mono_ns: int | None = None


@dataclass(frozen=True)
class RequestEvaluation:
    request_id: str
    passed: bool
    reasons: list[str]


def _rule_threshold_ms(rule: ThresholdRule, prompt_tokens: int | None) -> float:
    if rule.kind == "fixed_ms":
        assert rule.threshold_ms is not None
        return rule.threshold_ms
    if rule.kind == "linear_ms":
        assert rule.intercept_ms is not None and rule.per_token_ms is not None
        return float(rule.intercept_ms) + float(rule.per_token_ms) * float(prompt_tokens or 0)
    if rule.kind != "step_ms":
        raise ValueError(f"Unsupported threshold rule: {rule.kind}")
    prompt = float(prompt_tokens or 0)
    for bucket in rule.buckets:
        ceiling = bucket.get("max_input_tokens")
        if ceiling is None or prompt <= ceiling:
            return float(bucket["threshold_ms"])
    return float(rule.buckets[-1]["threshold_ms"])


def evaluate_request(outcome: RequestOutcome, slo: SloSpec) -> RequestEvaluation:
    reasons: list[str] = []
    if not outcome.success:
        reasons.append(outcome.error or "request_failed")
        return RequestEvaluation(request_id=outcome.request_id, passed=False, reasons=reasons)
    if slo.ttft_rule is not None:
        if outcome.ttft_ms is None:
            reasons.append("ttft_missing")
        else:
            threshold = _rule_threshold_ms(slo.ttft_rule, outcome.prompt_tokens)
            if outcome.ttft_ms > threshold:
                reasons.append(f"ttft_ms>{threshold}")
    if slo.tpot_rule is not None:
        if outcome.tpot_ms is None:
            reasons.append("tpot_missing")
        else:
            threshold = _rule_threshold_ms(slo.tpot_rule, outcome.prompt_tokens)
            if outcome.tpot_ms > threshold:
                reasons.append(f"tpot_ms>{threshold}")
    return RequestEvaluation(
        request_id=outcome.request_id,
        passed=not reasons,
        reasons=reasons,
    )


def summarize_evaluations(
    outcomes: list[RequestOutcome], slo: SloSpec
) -> tuple[list[RequestEvaluation], dict[str, Any]]:
    evaluations = [evaluate_request(item, slo) for item in outcomes]
    total = len(evaluations)
    passed = sum(1 for item in evaluations if item.passed)
    pass_rate = (passed / total) if total else 0.0
    return evaluations, {
        "request_count": total,
        "slo_pass_count": passed,
        "slo_pass_rate": pass_rate,
        "target_pass_rate": slo.target_pass_rate,
        "feasible": pass_rate >= slo.target_pass_rate,
    }
