"""Latency, throughput, token, CPU, GPU, and memory resource metrics.

Implements evaluation metrics for pipeline resource consumption and efficiency.
Supports per-document and per-stage breakdowns with percentile calculations.

Validates: Requirements 16.3, 16.4
"""
from __future__ import annotations

from dataclasses import dataclass

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Input Models
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class StageTimingRecord:
    """A single stage execution record with resource measurements.

    Captures timing, token usage, and hardware resource consumption
    for one processing stage of one document.
    """

    document_id: str
    stage_name: str
    start_time: float  # Unix timestamp (seconds)
    end_time: float  # Unix timestamp (seconds)
    input_tokens: int = 0
    output_tokens: int = 0
    gpu_memory_mb: float = 0.0
    cpu_seconds: float = 0.0
    gpu_seconds: float = 0.0

    @property
    def duration_seconds(self) -> float:
        """Wall-clock duration of this stage in seconds."""
        return self.end_time - self.start_time

    @property
    def total_tokens(self) -> int:
        """Sum of input and output tokens."""
        return self.input_tokens + self.output_tokens


# ---------------------------------------------------------------------------
# Percentile Helper
# ---------------------------------------------------------------------------


def compute_percentile(values: list[float], percentile: float) -> float:
    """Compute a percentile from a sorted list without numpy.

    Uses linear interpolation between nearest ranks.

    Args:
        values: List of numeric values (need not be pre-sorted).
        percentile: Percentile to compute (0-100).

    Returns:
        The interpolated percentile value.

    Raises:
        ValueError: If values is empty or percentile is out of range.
    """
    if not values:
        raise ValueError("Cannot compute percentile of empty list")
    if not (0.0 <= percentile <= 100.0):
        raise ValueError(f"Percentile must be between 0 and 100, got {percentile}")

    sorted_values = sorted(values)
    n = len(sorted_values)

    if n == 1:
        return sorted_values[0]

    # Compute the rank (0-indexed fractional position)
    rank = (percentile / 100.0) * (n - 1)
    lower_idx = int(rank)
    upper_idx = lower_idx + 1
    fraction = rank - lower_idx

    if upper_idx >= n:
        return sorted_values[-1]

    return sorted_values[lower_idx] + fraction * (
        sorted_values[upper_idx] - sorted_values[lower_idx]
    )


# ---------------------------------------------------------------------------
# Result Models
# ---------------------------------------------------------------------------


class LatencyPercentiles(BaseModel):
    """Latency percentile distribution in seconds."""

    p50: float = Field(ge=0.0)
    p90: float = Field(ge=0.0)
    p95: float = Field(ge=0.0)
    p99: float = Field(ge=0.0)
    mean: float = Field(ge=0.0)
    max: float = Field(ge=0.0)
    min: float = Field(ge=0.0)
    count: int = Field(ge=0)


class ThroughputMetrics(BaseModel):
    """Document throughput measurements."""

    documents_per_minute: float = Field(ge=0.0)
    documents_per_hour: float = Field(ge=0.0)
    total_documents: int = Field(ge=0)
    total_wall_seconds: float = Field(ge=0.0)


class TokenUsageMetrics(BaseModel):
    """Token consumption statistics."""

    total_input_tokens: int = Field(ge=0)
    total_output_tokens: int = Field(ge=0)
    total_tokens: int = Field(ge=0)
    mean_input_tokens_per_document: float = Field(ge=0.0)
    mean_output_tokens_per_document: float = Field(ge=0.0)
    mean_total_tokens_per_document: float = Field(ge=0.0)
    per_stage: dict[str, "StageTokenUsage"] = Field(default_factory=dict)


class StageTokenUsage(BaseModel):
    """Token usage breakdown for a single stage."""

    total_input_tokens: int = Field(ge=0)
    total_output_tokens: int = Field(ge=0)
    total_tokens: int = Field(ge=0)
    mean_input_tokens: float = Field(ge=0.0)
    mean_output_tokens: float = Field(ge=0.0)
    mean_total_tokens: float = Field(ge=0.0)
    count: int = Field(ge=0)


class CpuMetrics(BaseModel):
    """CPU resource consumption metrics."""

    total_cpu_seconds: float = Field(ge=0.0)
    mean_cpu_seconds_per_document: float = Field(ge=0.0)
    peak_cpu_seconds: float = Field(ge=0.0, description="Max CPU-seconds for a single document")


class GpuMetrics(BaseModel):
    """GPU resource consumption metrics."""

    total_gpu_seconds: float = Field(ge=0.0)
    mean_gpu_seconds_per_document: float = Field(ge=0.0)
    peak_gpu_memory_mb: float = Field(ge=0.0)
    mean_gpu_memory_mb: float = Field(ge=0.0)
    gpu_utilization_percent: float = Field(
        ge=0.0, le=100.0,
        description="Percentage of total wall time spent on GPU",
    )


class MemoryMetrics(BaseModel):
    """Memory consumption metrics."""

    peak_rss_memory_mb: float = Field(ge=0.0)
    mean_working_set_mb: float = Field(ge=0.0)


class EfficiencyMetrics(BaseModel):
    """Efficiency ratio metrics."""

    tokens_per_second: float = Field(ge=0.0)
    documents_per_gpu_second: float = Field(ge=0.0)
    fast_path_cpu_seconds: float = Field(ge=0.0)
    adjudication_cpu_seconds: float = Field(ge=0.0)
    fast_path_gpu_seconds: float = Field(ge=0.0)
    adjudication_gpu_seconds: float = Field(ge=0.0)
    fast_path_fraction: float = Field(
        ge=0.0, le=1.0,
        description="Fraction of total resource usage from fast-path stages",
    )
    adjudication_fraction: float = Field(
        ge=0.0, le=1.0,
        description="Fraction of total resource usage from adjudication stages",
    )


class StageLatencyBreakdown(BaseModel):
    """Per-stage latency statistics."""

    stage_name: str
    latency: LatencyPercentiles
    invocation_count: int = Field(ge=0)


class ResourceEvaluationReport(BaseModel):
    """Complete resource evaluation report."""

    latency: LatencyPercentiles
    per_stage_latency: list[StageLatencyBreakdown] = Field(default_factory=list)
    throughput: ThroughputMetrics
    token_usage: TokenUsageMetrics
    cpu: CpuMetrics
    gpu: GpuMetrics
    memory: MemoryMetrics
    efficiency: EfficiencyMetrics
    document_count: int = Field(ge=0)


# Rebuild model to resolve forward references
TokenUsageMetrics.model_rebuild()


# ---------------------------------------------------------------------------
# Computation Logic
# ---------------------------------------------------------------------------


# Stages considered as "adjudication" for resource split calculations
ADJUDICATION_STAGES: frozenset[str] = frozenset({
    "adjudication",
    "adjudicator",
    "9b_adjudication",
    "semantic_adjudication",
})


def _is_adjudication_stage(stage_name: str) -> bool:
    """Determine if a stage belongs to the adjudication path."""
    lower = stage_name.lower()
    return lower in ADJUDICATION_STAGES or "adjudicat" in lower


def _compute_latency_percentiles(durations: list[float]) -> LatencyPercentiles:
    """Compute latency percentile distribution from a list of durations."""
    if not durations:
        return LatencyPercentiles(
            p50=0.0, p90=0.0, p95=0.0, p99=0.0,
            mean=0.0, max=0.0, min=0.0, count=0,
        )

    return LatencyPercentiles(
        p50=compute_percentile(durations, 50.0),
        p90=compute_percentile(durations, 90.0),
        p95=compute_percentile(durations, 95.0),
        p99=compute_percentile(durations, 99.0),
        mean=sum(durations) / len(durations),
        max=max(durations),
        min=min(durations),
        count=len(durations),
    )


def _compute_document_durations(
    records: list[StageTimingRecord],
) -> dict[str, float]:
    """Compute total wall-clock duration per document.

    Uses min(start_time) to max(end_time) for each document to handle
    overlapping/parallel stages.
    """
    doc_times: dict[str, tuple[float, float]] = {}
    for r in records:
        if r.document_id not in doc_times:
            doc_times[r.document_id] = (r.start_time, r.end_time)
        else:
            existing = doc_times[r.document_id]
            doc_times[r.document_id] = (
                min(existing[0], r.start_time),
                max(existing[1], r.end_time),
            )
    return {doc_id: end - start for doc_id, (start, end) in doc_times.items()}


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def compute_latency_metrics(
    records: list[StageTimingRecord],
) -> tuple[LatencyPercentiles, list[StageLatencyBreakdown]]:
    """Compute per-document and per-stage latency metrics.

    Per-document latency is the wall-clock time from the earliest stage
    start to the latest stage end for each document.

    Args:
        records: Stage timing records.

    Returns:
        Tuple of (overall document latency, per-stage breakdown).
    """
    if not records:
        return (
            LatencyPercentiles(
                p50=0.0, p90=0.0, p95=0.0, p99=0.0,
                mean=0.0, max=0.0, min=0.0, count=0,
            ),
            [],
        )

    # Per-document latency
    doc_durations = _compute_document_durations(records)
    overall = _compute_latency_percentiles(list(doc_durations.values()))

    # Per-stage latency
    stage_durations: dict[str, list[float]] = {}
    for r in records:
        stage_durations.setdefault(r.stage_name, []).append(r.duration_seconds)

    per_stage = [
        StageLatencyBreakdown(
            stage_name=stage,
            latency=_compute_latency_percentiles(durations),
            invocation_count=len(durations),
        )
        for stage, durations in sorted(stage_durations.items())
    ]

    return overall, per_stage


def compute_throughput_metrics(
    records: list[StageTimingRecord],
) -> ThroughputMetrics:
    """Compute document throughput from timing records.

    Args:
        records: Stage timing records.

    Returns:
        ThroughputMetrics with documents/minute and documents/hour.
    """
    if not records:
        return ThroughputMetrics(
            documents_per_minute=0.0,
            documents_per_hour=0.0,
            total_documents=0,
            total_wall_seconds=0.0,
        )

    doc_ids = {r.document_id for r in records}
    total_docs = len(doc_ids)

    # Total wall time: earliest start to latest end across all records
    earliest = min(r.start_time for r in records)
    latest = max(r.end_time for r in records)
    total_wall = latest - earliest

    if total_wall <= 0.0:
        return ThroughputMetrics(
            documents_per_minute=0.0,
            documents_per_hour=0.0,
            total_documents=total_docs,
            total_wall_seconds=0.0,
        )

    docs_per_second = total_docs / total_wall

    return ThroughputMetrics(
        documents_per_minute=docs_per_second * 60.0,
        documents_per_hour=docs_per_second * 3600.0,
        total_documents=total_docs,
        total_wall_seconds=total_wall,
    )


def compute_token_usage_metrics(
    records: list[StageTimingRecord],
) -> TokenUsageMetrics:
    """Compute token usage statistics per document and per stage.

    Args:
        records: Stage timing records.

    Returns:
        TokenUsageMetrics with aggregate and per-stage breakdowns.
    """
    if not records:
        return TokenUsageMetrics(
            total_input_tokens=0,
            total_output_tokens=0,
            total_tokens=0,
            mean_input_tokens_per_document=0.0,
            mean_output_tokens_per_document=0.0,
            mean_total_tokens_per_document=0.0,
            per_stage={},
        )

    total_input = sum(r.input_tokens for r in records)
    total_output = sum(r.output_tokens for r in records)
    total = total_input + total_output

    doc_ids = {r.document_id for r in records}
    n_docs = len(doc_ids)

    # Per-stage breakdown
    stage_records: dict[str, list[StageTimingRecord]] = {}
    for r in records:
        stage_records.setdefault(r.stage_name, []).append(r)

    per_stage: dict[str, StageTokenUsage] = {}
    for stage, stage_recs in sorted(stage_records.items()):
        s_input = sum(r.input_tokens for r in stage_recs)
        s_output = sum(r.output_tokens for r in stage_recs)
        s_total = s_input + s_output
        count = len(stage_recs)
        per_stage[stage] = StageTokenUsage(
            total_input_tokens=s_input,
            total_output_tokens=s_output,
            total_tokens=s_total,
            mean_input_tokens=s_input / count if count > 0 else 0.0,
            mean_output_tokens=s_output / count if count > 0 else 0.0,
            mean_total_tokens=s_total / count if count > 0 else 0.0,
            count=count,
        )

    return TokenUsageMetrics(
        total_input_tokens=total_input,
        total_output_tokens=total_output,
        total_tokens=total,
        mean_input_tokens_per_document=total_input / n_docs if n_docs > 0 else 0.0,
        mean_output_tokens_per_document=total_output / n_docs if n_docs > 0 else 0.0,
        mean_total_tokens_per_document=total / n_docs if n_docs > 0 else 0.0,
        per_stage=per_stage,
    )


def compute_cpu_metrics(
    records: list[StageTimingRecord],
) -> CpuMetrics:
    """Compute CPU resource consumption metrics.

    Args:
        records: Stage timing records.

    Returns:
        CpuMetrics with totals and per-document statistics.
    """
    if not records:
        return CpuMetrics(
            total_cpu_seconds=0.0,
            mean_cpu_seconds_per_document=0.0,
            peak_cpu_seconds=0.0,
        )

    total_cpu = sum(r.cpu_seconds for r in records)

    # Per-document CPU totals
    doc_cpu: dict[str, float] = {}
    for r in records:
        doc_cpu[r.document_id] = doc_cpu.get(r.document_id, 0.0) + r.cpu_seconds

    n_docs = len(doc_cpu)
    peak = max(doc_cpu.values()) if doc_cpu else 0.0

    return CpuMetrics(
        total_cpu_seconds=total_cpu,
        mean_cpu_seconds_per_document=total_cpu / n_docs if n_docs > 0 else 0.0,
        peak_cpu_seconds=peak,
    )


def compute_gpu_metrics(
    records: list[StageTimingRecord],
) -> GpuMetrics:
    """Compute GPU resource consumption metrics.

    Args:
        records: Stage timing records.

    Returns:
        GpuMetrics with totals, peaks, and utilization percentage.
    """
    if not records:
        return GpuMetrics(
            total_gpu_seconds=0.0,
            mean_gpu_seconds_per_document=0.0,
            peak_gpu_memory_mb=0.0,
            mean_gpu_memory_mb=0.0,
            gpu_utilization_percent=0.0,
        )

    total_gpu = sum(r.gpu_seconds for r in records)

    # Per-document GPU totals
    doc_gpu: dict[str, float] = {}
    for r in records:
        doc_gpu[r.document_id] = doc_gpu.get(r.document_id, 0.0) + r.gpu_seconds

    n_docs = len(doc_gpu)

    # GPU memory stats
    gpu_mem_values = [r.gpu_memory_mb for r in records if r.gpu_memory_mb > 0.0]
    peak_gpu_mem = max(gpu_mem_values) if gpu_mem_values else 0.0
    mean_gpu_mem = (
        sum(gpu_mem_values) / len(gpu_mem_values) if gpu_mem_values else 0.0
    )

    # GPU utilization: fraction of wall time spent on GPU work
    earliest = min(r.start_time for r in records)
    latest = max(r.end_time for r in records)
    total_wall = latest - earliest

    utilization = (
        (total_gpu / total_wall) * 100.0 if total_wall > 0.0 else 0.0
    )
    # Cap at 100% (parallel GPU stages could theoretically exceed wall time)
    utilization = min(utilization, 100.0)

    return GpuMetrics(
        total_gpu_seconds=total_gpu,
        mean_gpu_seconds_per_document=total_gpu / n_docs if n_docs > 0 else 0.0,
        peak_gpu_memory_mb=peak_gpu_mem,
        mean_gpu_memory_mb=mean_gpu_mem,
        gpu_utilization_percent=utilization,
    )


def compute_memory_metrics(
    records: list[StageTimingRecord],
    rss_samples_mb: list[float] | None = None,
) -> MemoryMetrics:
    """Compute memory consumption metrics.

    Uses gpu_memory_mb as a proxy for working set if no explicit RSS
    samples are provided. When rss_samples_mb is given, it takes
    precedence for peak and mean calculations.

    Args:
        records: Stage timing records.
        rss_samples_mb: Optional explicit RSS memory samples in MB.

    Returns:
        MemoryMetrics with peak and mean working set.
    """
    if rss_samples_mb:
        return MemoryMetrics(
            peak_rss_memory_mb=max(rss_samples_mb),
            mean_working_set_mb=sum(rss_samples_mb) / len(rss_samples_mb),
        )

    if not records:
        return MemoryMetrics(peak_rss_memory_mb=0.0, mean_working_set_mb=0.0)

    # Use gpu_memory_mb as working set proxy
    mem_values = [r.gpu_memory_mb for r in records if r.gpu_memory_mb > 0.0]
    if not mem_values:
        return MemoryMetrics(peak_rss_memory_mb=0.0, mean_working_set_mb=0.0)

    return MemoryMetrics(
        peak_rss_memory_mb=max(mem_values),
        mean_working_set_mb=sum(mem_values) / len(mem_values),
    )


def compute_efficiency_metrics(
    records: list[StageTimingRecord],
) -> EfficiencyMetrics:
    """Compute efficiency ratios including tokens/second and resource splits.

    Fast-path vs adjudication split is determined by stage name matching.

    Args:
        records: Stage timing records.

    Returns:
        EfficiencyMetrics with ratios and resource splits.
    """
    if not records:
        return EfficiencyMetrics(
            tokens_per_second=0.0,
            documents_per_gpu_second=0.0,
            fast_path_cpu_seconds=0.0,
            adjudication_cpu_seconds=0.0,
            fast_path_gpu_seconds=0.0,
            adjudication_gpu_seconds=0.0,
            fast_path_fraction=0.0,
            adjudication_fraction=0.0,
        )

    total_tokens = sum(r.total_tokens for r in records)
    total_wall = max(r.end_time for r in records) - min(r.start_time for r in records)
    total_gpu = sum(r.gpu_seconds for r in records)
    n_docs = len({r.document_id for r in records})

    tokens_per_second = total_tokens / total_wall if total_wall > 0.0 else 0.0
    docs_per_gpu_second = n_docs / total_gpu if total_gpu > 0.0 else 0.0

    # Resource split
    fast_cpu = 0.0
    adj_cpu = 0.0
    fast_gpu = 0.0
    adj_gpu = 0.0

    for r in records:
        if _is_adjudication_stage(r.stage_name):
            adj_cpu += r.cpu_seconds
            adj_gpu += r.gpu_seconds
        else:
            fast_cpu += r.cpu_seconds
            fast_gpu += r.gpu_seconds

    total_resource = fast_cpu + adj_cpu + fast_gpu + adj_gpu
    fast_total = fast_cpu + fast_gpu
    adj_total = adj_cpu + adj_gpu

    fast_fraction = fast_total / total_resource if total_resource > 0.0 else 0.0
    adj_fraction = adj_total / total_resource if total_resource > 0.0 else 0.0

    return EfficiencyMetrics(
        tokens_per_second=tokens_per_second,
        documents_per_gpu_second=docs_per_gpu_second,
        fast_path_cpu_seconds=fast_cpu,
        adjudication_cpu_seconds=adj_cpu,
        fast_path_gpu_seconds=fast_gpu,
        adjudication_gpu_seconds=adj_gpu,
        fast_path_fraction=fast_fraction,
        adjudication_fraction=adj_fraction,
    )


def evaluate_resources(
    records: list[StageTimingRecord],
    rss_samples_mb: list[float] | None = None,
) -> ResourceEvaluationReport:
    """Run full resource evaluation producing a complete report.

    Args:
        records: List of stage timing records from pipeline execution.
        rss_samples_mb: Optional explicit RSS memory samples.

    Returns:
        ResourceEvaluationReport with all resource metrics.
    """
    latency, per_stage_latency = compute_latency_metrics(records)
    throughput = compute_throughput_metrics(records)
    token_usage = compute_token_usage_metrics(records)
    cpu = compute_cpu_metrics(records)
    gpu = compute_gpu_metrics(records)
    memory = compute_memory_metrics(records, rss_samples_mb)
    efficiency = compute_efficiency_metrics(records)

    doc_count = len({r.document_id for r in records}) if records else 0

    return ResourceEvaluationReport(
        latency=latency,
        per_stage_latency=per_stage_latency,
        throughput=throughput,
        token_usage=token_usage,
        cpu=cpu,
        gpu=gpu,
        memory=memory,
        efficiency=efficiency,
        document_count=doc_count,
    )