from __future__ import annotations

import json
import math
import statistics
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Sequence

from .trace import TraceRequest, WindowRecord

if TYPE_CHECKING:
    from .spec import StudySpec


EPSILON = 1e-9

FEATURE_NAMES = [
    "L.log_mean_length",
    "L.log_p95_over_mean_length",
    "L.cv_length",
    "C.log_mean_hit_length",
    "C.log_p95_over_mean_hit_length",
    "C.cv_hit_length",
    "C.hit_rate",
    "A.log_request_rate_per_gpu",
    "A.cv_interarrival",
    "A.log_fano_1s",
]

FAMILY_SLICES = {
    "L": slice(0, 3),
    "C": slice(3, 7),
    "A": slice(7, 10),
}

LENGTH_MODES = {"total", "input", "output"}


@dataclass(frozen=True)
class WorkloadProfile:
    window_id: str
    trace_type: str
    request_count: int
    duration_s: float
    gpu_count: int
    length_mode: str
    feature_names: list[str]
    vector: list[float]
    stats: dict[str, Any]

    def to_dict(self) -> dict[str, Any]:
        return {
            "window_id": self.window_id,
            "trace_type": self.trace_type,
            "request_count": self.request_count,
            "duration_s": self.duration_s,
            "gpu_count": self.gpu_count,
            "length_mode": self.length_mode,
            "feature_names": self.feature_names,
            "vector": self.vector,
            "stats": self.stats,
        }


@dataclass(frozen=True)
class RobustScale:
    feature_names: list[str]
    center: list[float]
    scale: list[float]

    def transform(self, vector: Sequence[float]) -> list[float]:
        return [
            (float(value) - self.center[idx]) / self.scale[idx]
            for idx, value in enumerate(vector)
        ]

    def to_dict(self) -> dict[str, Any]:
        return {
            "feature_names": self.feature_names,
            "center": self.center,
            "scale": self.scale,
        }


def resolve_length_mode(*, request_mode: str | None = None, length_mode: str = "auto") -> str:
    normalized = str(length_mode or "auto").strip().lower()
    if normalized == "auto":
        return (
            "output"
            if str(request_mode or "").strip().lower() == "decode_only"
            else "total"
        )
    if normalized not in LENGTH_MODES:
        raise ValueError(
            "length_mode must be one of: auto, total, input, output."
        )
    return normalized


def build_workload_profile(
    requests: list[TraceRequest],
    window: WindowRecord,
    *,
    gpu_count: int,
    length_mode: str = "total",
) -> WorkloadProfile:
    if gpu_count <= 0:
        raise ValueError("gpu_count must be > 0.")
    if length_mode not in LENGTH_MODES:
        raise ValueError(f"Unsupported length_mode: {length_mode}")

    duration_s = _duration_s(requests, window)
    input_lengths = [float(item.prompt_tokens_hint or 0) for item in requests]
    output_lengths = [float(item.completion_tokens_hint or 0) for item in requests]
    profile_lengths = [
        _profile_length(input_len, output_len, length_mode=length_mode)
        for input_len, output_len in zip(input_lengths, output_lengths)
    ]
    hit_lengths, cache_stats = _ideal_cache_hit_lengths(
        requests,
        input_lengths=input_lengths,
        block_size=_block_size(window),
    )
    arrival_stats = _arrival_stats(requests, duration_s=duration_s, gpu_count=gpu_count)

    length_stats = _series_stats(profile_lengths)
    hit_stats = _series_stats(hit_lengths)
    total_profile_length = sum(profile_lengths)
    total_input_length = sum(input_lengths)
    total_hit_length = sum(hit_lengths)
    feature_hit_rate = (
        float(total_hit_length / max(total_profile_length, EPSILON))
        if total_profile_length > 0
        else 0.0
    )
    input_hit_rate = (
        float(total_hit_length / max(total_input_length, EPSILON))
        if total_input_length > 0
        else 0.0
    )

    vector = [
        math.log1p(length_stats["mean"]),
        math.log1p(length_stats["p95"] / max(length_stats["mean"], EPSILON)),
        length_stats["cv"],
        math.log1p(hit_stats["mean"]),
        math.log1p(hit_stats["p95"] / max(hit_stats["mean"], EPSILON)),
        hit_stats["cv"],
        feature_hit_rate,
        math.log1p(arrival_stats["request_rate_per_gpu"]),
        arrival_stats["interarrival_cv"],
        math.log1p(arrival_stats["fano_1s"]),
    ]

    return WorkloadProfile(
        window_id=window.window_id,
        trace_type=window.trace_type,
        request_count=len(requests),
        duration_s=duration_s,
        gpu_count=int(gpu_count),
        length_mode=length_mode,
        feature_names=list(FEATURE_NAMES),
        vector=[float(item) for item in vector],
        stats={
            "length": {
                **length_stats,
                "mode": length_mode,
                "total": total_profile_length,
                "input_total": total_input_length,
                "output_total": sum(output_lengths),
            },
            "cache": {
                **hit_stats,
                **cache_stats,
                "total_hit_length": total_hit_length,
                "hit_rate": feature_hit_rate,
                "input_hit_rate": input_hit_rate,
            },
            "arrival": arrival_stats,
        },
    )


def build_study_workload_profile(
    study: "StudySpec",
    requests: list[TraceRequest],
    window: WindowRecord,
) -> WorkloadProfile:
    """Canonical L-C-A profile for a study's loaded window.

    This is the single source of truth for the paper's 10-dimensional L-C-A
    feature vector used by the harness prompt and (later) by Stop-A.
    """
    mode = resolve_length_mode(
        request_mode=study.trace.request_mode,
        length_mode="auto",
    )
    return build_workload_profile(
        requests,
        window,
        gpu_count=study.hardware.gpu_count,
        length_mode=mode,
    )


def fit_robust_scale(profiles: Sequence[WorkloadProfile]) -> RobustScale:
    if not profiles:
        raise ValueError("At least one profile is required to fit a robust scale.")
    centers: list[float] = []
    scales: list[float] = []
    for idx in range(len(FEATURE_NAMES)):
        values = [float(profile.vector[idx]) for profile in profiles]
        median = _percentile(values, 50.0)
        iqr = _percentile(values, 75.0) - _percentile(values, 25.0)
        centers.append(float(median))
        scales.append(float(iqr if abs(iqr) > EPSILON else 1.0))
    return RobustScale(feature_names=list(FEATURE_NAMES), center=centers, scale=scales)


def profile_similarity(
    left: WorkloadProfile,
    right: WorkloadProfile,
    *,
    scale: RobustScale | None = None,
) -> float:
    scaler = scale or fit_robust_scale([left, right])
    z_left = scaler.transform(left.vector)
    z_right = scaler.transform(right.vector)
    return _similarity_from_z(z_left, z_right)


def similarity_report(profiles: Sequence[WorkloadProfile]) -> dict[str, Any]:
    if not profiles:
        raise ValueError("At least one profile is required.")
    scale = fit_robust_scale(profiles)
    transformed = [scale.transform(profile.vector) for profile in profiles]
    rows: list[dict[str, Any]] = []
    matrix: list[list[float]] = []
    for i, left in enumerate(profiles):
        row_values: list[float] = []
        for j, right in enumerate(profiles):
            sim = _similarity_from_z(transformed[i], transformed[j])
            row_values.append(sim)
            rows.append(
                {
                    "left": left.window_id,
                    "right": right.window_id,
                    "similarity": sim,
                    "family_similarity": _family_similarity(transformed[i], transformed[j]),
                }
            )
        matrix.append(row_values)
    return {
        "feature_names": list(FEATURE_NAMES),
        "scaler": scale.to_dict(),
        "windows": [profile.window_id for profile in profiles],
        "matrix": matrix,
        "pairs": rows,
    }


@dataclass(frozen=True)
class ConvergencePoint:
    converged: bool
    stop_index: int
    stop_time_s: float
    fraction: float
    family_similarity: dict[str, float]
    checks: list[dict[str, Any]]

    def to_dict(self) -> dict[str, Any]:
        return {
            "converged": self.converged,
            "stop_index": self.stop_index,
            "stop_time_s": self.stop_time_s,
            "fraction": self.fraction,
            "family_similarity": self.family_similarity,
            "checks": self.checks,
        }


def find_convergence_prefix(
    requests: list[TraceRequest],
    window: WindowRecord,
    *,
    gpu_count: int,
    length_mode: str = "total",
    tau: float = 0.9,
    tau_c: float = 0.92,
    stable_checks: int = 3,
    max_checks: int = 20,
    min_fraction: float = 0.1,
) -> ConvergencePoint:
    """Earliest arrival-ordered prefix whose offered L-C-A converges to the full set.

    The L-C-A vector is a deterministic function of the trace metadata, so the
    convergence of prefix-vs-full is itself deterministic (the paper's Fig. 9
    curve). Stop-A replays only up to this prefix. A prefix counts as converged
    when the L and A family similarities reach ``tau`` and the (slowest) C family
    similarity reaches the stricter ``tau_c`` for ``stable_checks`` consecutive
    checkpoints. If that never happens within the window the point reports the
    full set (converged=False), which keeps the C-gate honest: an unconverged C
    means the probe must replay the whole window rather than stop early.
    """
    total = len(requests)
    if total == 0:
        return ConvergencePoint(
            converged=False,
            stop_index=0,
            stop_time_s=0.0,
            fraction=1.0,
            family_similarity={"L": 1.0, "C": 1.0, "A": 1.0},
            checks=[],
        )
    # Compare each arrival-ordered prefix to the whole set, both measured over
    # their own elapsed span so the A (rate) dimension is comparable rather than
    # diluted by the fixed window length.
    target = _prefix_profile(
        requests, total, window, gpu_count=gpu_count, length_mode=length_mode
    )
    indices = _checkpoint_indices(
        total, max_checks=max_checks, min_fraction=min_fraction
    )
    checks: list[dict[str, Any]] = []
    consecutive = 0
    converged_index: int | None = None
    converged_sims: dict[str, float] | None = None
    for index in indices:
        prefix = _prefix_profile(
            requests, index, window, gpu_count=gpu_count, length_mode=length_mode
        )
        sims = _family_similarity(target.vector, prefix.vector)
        checks.append(
            {
                "index": index,
                "fraction": float(index / total),
                "time_s": float(requests[index - 1].arrival_s),
                "family_similarity": sims,
            }
        )
        passed = sims["L"] >= tau and sims["A"] >= tau and sims["C"] >= tau_c
        consecutive = consecutive + 1 if passed else 0
        if consecutive >= stable_checks and converged_index is None:
            converged_index = index
            converged_sims = sims
            break
    if converged_index is None:
        last_sims = checks[-1]["family_similarity"] if checks else {"L": 1.0, "C": 1.0, "A": 1.0}
        return ConvergencePoint(
            converged=False,
            stop_index=total,
            stop_time_s=float(requests[-1].arrival_s),
            fraction=1.0,
            family_similarity=last_sims,
            checks=checks,
        )
    return ConvergencePoint(
        converged=True,
        stop_index=converged_index,
        stop_time_s=float(requests[converged_index - 1].arrival_s),
        fraction=float(converged_index / total),
        family_similarity=converged_sims or {},
        checks=checks,
    )


def _prefix_profile(
    requests: list[TraceRequest],
    index: int,
    window: WindowRecord,
    *,
    gpu_count: int,
    length_mode: str,
) -> WorkloadProfile:
    prefix = requests[:index]
    start = float(prefix[0].arrival_s) if prefix else float(window.window_start)
    end = float(prefix[-1].arrival_s) if prefix else float(window.window_start)
    prefix_window = WindowRecord(
        window_id=window.window_id,
        trace_path=window.trace_path,
        trace_type=window.trace_type,
        window_start=start,
        window_end=end,
        source_payload=window.source_payload,
    )
    return build_workload_profile(
        prefix, prefix_window, gpu_count=gpu_count, length_mode=length_mode
    )


def _checkpoint_indices(total: int, *, max_checks: int, min_fraction: float) -> list[int]:
    start = max(1, int(math.ceil(min_fraction * total)))
    if total <= max_checks:
        candidates = range(start, total + 1)
    else:
        step = max(1, total // max_checks)
        candidates = list(range(start, total + 1, step))
        if candidates and candidates[-1] != total:
            candidates.append(total)
    seen: list[int] = []
    for value in candidates:
        clamped = min(total, max(1, int(value)))
        if not seen or seen[-1] != clamped:
            seen.append(clamped)
    return seen


def dumps_profile(profile: WorkloadProfile) -> str:
    return json.dumps(profile.to_dict(), ensure_ascii=False, indent=2) + "\n"


def _duration_s(requests: list[TraceRequest], window: WindowRecord) -> float:
    duration = max(float(window.window_end) - float(window.window_start), 0.0)
    if duration > 0:
        return duration
    if len(requests) >= 2:
        return max(0.0, float(requests[-1].arrival_s) - float(requests[0].arrival_s))
    return 0.0


def _profile_length(input_length: float, output_length: float, *, length_mode: str) -> float:
    if length_mode == "input":
        return max(input_length, 0.0)
    if length_mode == "output":
        return max(output_length, 0.0)
    return max(input_length, 0.0) + max(output_length, 0.0)


def _block_size(window: WindowRecord) -> int:
    value = window.source_payload.get("block_size")
    if isinstance(value, bool):
        return 1
    if isinstance(value, (int, float)) and value > 0:
        return int(value)
    if isinstance(value, str) and value.strip():
        try:
            parsed = int(value)
        except ValueError:
            return 1
        return parsed if parsed > 0 else 1
    return 1


def _ideal_cache_hit_lengths(
    requests: list[TraceRequest],
    *,
    input_lengths: list[float],
    block_size: int,
) -> tuple[list[float], dict[str, Any]]:
    seen_hashes: set[Any] = set()
    hit_lengths: list[float] = []
    total_blocks = 0
    repeated_blocks = 0
    rows_with_hash_ids = 0
    for request, input_length in zip(requests, input_lengths):
        hash_ids = request.metadata.get("hash_ids")
        if not isinstance(hash_ids, list):
            hit_lengths.append(0.0)
            continue
        rows_with_hash_ids += 1
        repeated_for_request = 0
        for hash_id in hash_ids:
            total_blocks += 1
            if hash_id in seen_hashes:
                repeated_blocks += 1
                repeated_for_request += 1
            else:
                seen_hashes.add(hash_id)
        hit_lengths.append(float(min(max(input_length, 0.0), repeated_for_request * block_size)))
    return hit_lengths, {
        "block_size": block_size,
        "rows_with_hash_ids": rows_with_hash_ids,
        "total_blocks": total_blocks,
        "repeated_blocks": repeated_blocks,
        "repeated_block_ratio": (
            float(repeated_blocks / total_blocks) if total_blocks else 0.0
        ),
    }


def _arrival_stats(
    requests: list[TraceRequest],
    *,
    duration_s: float,
    gpu_count: int,
) -> dict[str, Any]:
    arrivals = [float(item.arrival_s) for item in requests]
    interarrivals = [
        max(0.0, arrivals[idx] - arrivals[idx - 1])
        for idx in range(1, len(arrivals))
    ]
    per_second_counts = _per_second_counts(arrivals, duration_s=duration_s)
    qps = float(len(requests) / duration_s) if duration_s > 0 else 0.0
    return {
        "request_rate": qps,
        "request_rate_per_gpu": float(qps / gpu_count) if gpu_count > 0 else 0.0,
        "interarrival_cv": _cv(interarrivals),
        "fano_1s": _fano(per_second_counts),
        "one_second_count_mean": statistics.fmean(per_second_counts)
        if per_second_counts
        else 0.0,
        "one_second_count_variance": statistics.pvariance(per_second_counts)
        if len(per_second_counts) >= 2
        else 0.0,
        "one_second_bin_count": len(per_second_counts),
    }


def _per_second_counts(arrivals: list[float], *, duration_s: float) -> list[float]:
    if duration_s <= 0:
        return [float(len(arrivals))] if arrivals else []
    bin_count = max(1, int(math.ceil(duration_s)))
    counts = [0.0 for _ in range(bin_count)]
    for arrival in arrivals:
        if arrival < 0:
            continue
        idx = int(math.floor(arrival))
        if 0 <= idx < bin_count:
            counts[idx] += 1.0
    return counts


def _series_stats(values: list[float]) -> dict[str, float]:
    return {
        "count": float(len(values)),
        "mean": statistics.fmean(values) if values else 0.0,
        "p50": _percentile(values, 50.0),
        "p95": _percentile(values, 95.0),
        "cv": _cv(values),
    }


def _cv(values: list[float]) -> float:
    if not values:
        return 0.0
    mean = statistics.fmean(values)
    if abs(mean) <= EPSILON:
        return 0.0
    return float(statistics.pstdev(values) / mean) if len(values) >= 2 else 0.0


def _fano(values: list[float]) -> float:
    if not values:
        return 0.0
    mean = statistics.fmean(values)
    if abs(mean) <= EPSILON:
        return 0.0
    return float(statistics.pvariance(values) / mean) if len(values) >= 2 else 0.0


def _percentile(values: Sequence[float], p: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(float(item) for item in values)
    if len(ordered) == 1:
        return ordered[0]
    rank = (p / 100.0) * (len(ordered) - 1)
    lower = int(math.floor(rank))
    upper = int(math.ceil(rank))
    if lower == upper:
        return ordered[lower]
    weight = rank - lower
    return float(ordered[lower] * (1.0 - weight) + ordered[upper] * weight)


def _similarity_from_z(left: Sequence[float], right: Sequence[float]) -> float:
    distance = math.sqrt(
        sum((float(lval) - float(rval)) ** 2 for lval, rval in zip(left, right))
    )
    return float(math.exp(-distance))


def _family_similarity(left: Sequence[float], right: Sequence[float]) -> dict[str, float]:
    result: dict[str, float] = {}
    for family, family_slice in FAMILY_SLICES.items():
        result[family] = _similarity_from_z(left[family_slice], right[family_slice])
    return result
