from __future__ import annotations

import json
import statistics
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable

from agentic_pd_hybrid.trace import TraceRequest, load_trace


@dataclass(frozen=True)
class ProfileConfig:
    trace_path: Path
    output_path: Path | None = None
    metrics_path: Path | None = None
    baseline_metrics_path: Path | None = None
    candidate_metrics_path: Path | None = None
    direct_max_uncached_tokens: int = 2048


def write_profile(config: ProfileConfig) -> dict[str, Any]:
    report = build_profile(config)
    if config.output_path is not None:
        config.output_path.parent.mkdir(parents=True, exist_ok=True)
        with config.output_path.open("w", encoding="utf-8") as handle:
            json.dump(report, handle, indent=2, sort_keys=True)
    return report


def build_profile(config: ProfileConfig) -> dict[str, Any]:
    requests = load_trace(config.trace_path)
    features = _build_trace_features(
        requests,
        direct_max_uncached_tokens=config.direct_max_uncached_tokens,
    )
    report: dict[str, Any] = {
        "trace_path": str(config.trace_path),
        "direct_max_uncached_tokens": config.direct_max_uncached_tokens,
        "trace_profile": _trace_profile(requests, features),
    }

    if config.metrics_path is not None:
        metrics = _load_jsonl(config.metrics_path)
        report["metrics_path"] = str(config.metrics_path)
        report["metrics_profile"] = _metrics_profile(metrics, features)

    if (
        config.baseline_metrics_path is not None
        and config.candidate_metrics_path is not None
    ):
        baseline = _load_jsonl(config.baseline_metrics_path)
        candidate = _load_jsonl(config.candidate_metrics_path)
        report["baseline_metrics_path"] = str(config.baseline_metrics_path)
        report["candidate_metrics_path"] = str(config.candidate_metrics_path)
        report["baseline_profile"] = _metrics_profile(baseline, features)
        report["candidate_profile"] = _metrics_profile(candidate, features)
        report["paired_comparison"] = _paired_comparison(
            baseline=baseline,
            candidate=candidate,
            features=features,
        )

    return report


def print_profile_summary(report: dict[str, Any]) -> None:
    trace = report["trace_profile"]
    print(
        "trace: "
        f"{trace['request_count']} requests, "
        f"{trace['session_count']} sessions, "
        f"{trace['multi_turn_session_count']} multi-turn sessions"
    )
    print(
        "direct-eligible turns: "
        f"{trace['direct_eligible_turn2plus_count']}/"
        f"{trace['turn2plus_count']} "
        f"({trace['direct_eligible_turn2plus_ratio']:.3f})"
    )
    append_stats = trace.get("append_input_tokens_stats")
    output_stats = trace.get("output_tokens_stats")
    if append_stats is not None:
        print(
            "append tokens: "
            f"mean={append_stats['mean']:.1f} "
            f"p50={append_stats['p50']:.1f} "
            f"p90={append_stats['p90']:.1f} "
            f"p99={append_stats['p99']:.1f}"
        )
    if output_stats is not None:
        print(
            "output tokens: "
            f"mean={output_stats['mean']:.1f} "
            f"p50={output_stats['p50']:.1f} "
            f"p90={output_stats['p90']:.1f} "
            f"p99={output_stats['p99']:.1f}"
        )

    comparison = report.get("paired_comparison")
    if isinstance(comparison, dict):
        overall = comparison.get("overall", {})
        delta = overall.get("latency_delta_s_stats")
        if delta is not None:
            print(
                "candidate - baseline E2E: "
                f"mean={delta['mean']:.3f}s "
                f"p50={delta['p50']:.3f}s "
                f"p90={delta['p90']:.3f}s"
            )
        print(
            "paired wins/losses: "
            f"{overall.get('candidate_faster_count', 0)} faster, "
            f"{overall.get('candidate_slower_count', 0)} slower, "
            f"{overall.get('paired_count', 0)} paired"
        )


@dataclass(frozen=True)
class _TraceFeature:
    request_id: str
    session_id: str
    turn_id: int
    input_length: int
    output_length: int
    resident_tokens: int
    append_input_tokens: int | None
    inter_turn_gap_s: float | None
    overlap_blocks_with_previous: int | None
    overlap_ratio_with_previous: float | None
    direct_eligible: bool
    turn_class: str
    append_bin: str
    input_bin: str
    output_bin: str
    resident_bin: str


def _build_trace_features(
    requests: list[TraceRequest],
    *,
    direct_max_uncached_tokens: int,
) -> dict[str, _TraceFeature]:
    ordered_by_session: dict[str, list[TraceRequest]] = defaultdict(list)
    for request in requests:
        ordered_by_session[request.session_id].append(request)

    previous_by_request_id: dict[str, TraceRequest | None] = {}
    for session_requests in ordered_by_session.values():
        ordered = sorted(
            session_requests,
            key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
        )
        previous: TraceRequest | None = None
        for request in ordered:
            previous_by_request_id[request.request_id] = previous
            previous = request

    features: dict[str, _TraceFeature] = {}
    for request in requests:
        previous = previous_by_request_id.get(request.request_id)
        append_input_tokens: int | None = None
        inter_turn_gap_s: float | None = None
        overlap_blocks: int | None = None
        overlap_ratio: float | None = None
        direct_eligible = False

        if previous is not None:
            append_input_tokens = request.input_length - (
                previous.input_length + previous.output_length
            )
            inter_turn_gap_s = request.timestamp_s - previous.timestamp_s
            previous_blocks = set(previous.hash_ids)
            overlap_blocks = sum(1 for block in request.hash_ids if block in previous_blocks)
            overlap_ratio = (
                overlap_blocks / len(request.hash_ids) if request.hash_ids else 0.0
            )
            direct_eligible = (
                append_input_tokens > 0
                and append_input_tokens <= direct_max_uncached_tokens
                and overlap_blocks > 0
            )

        features[request.request_id] = _TraceFeature(
            request_id=request.request_id,
            session_id=request.session_id,
            turn_id=request.turn_id,
            input_length=request.input_length,
            output_length=request.output_length,
            resident_tokens=request.input_length + request.output_length,
            append_input_tokens=append_input_tokens,
            inter_turn_gap_s=inter_turn_gap_s,
            overlap_blocks_with_previous=overlap_blocks,
            overlap_ratio_with_previous=overlap_ratio,
            direct_eligible=direct_eligible,
            turn_class="turn1" if request.turn_id <= 1 else "turn2plus",
            append_bin=_token_bin(append_input_tokens),
            input_bin=_token_bin(request.input_length),
            output_bin=_token_bin(request.output_length),
            resident_bin=_token_bin(request.input_length + request.output_length),
        )
    return features


def _trace_profile(
    requests: list[TraceRequest],
    features: dict[str, _TraceFeature],
) -> dict[str, Any]:
    session_turns = Counter(request.session_id for request in requests)
    turn2plus = [feature for feature in features.values() if feature.turn_id > 1]
    direct_eligible = [feature for feature in turn2plus if feature.direct_eligible]
    append_values = [
        feature.append_input_tokens
        for feature in turn2plus
        if feature.append_input_tokens is not None
    ]
    positive_append_values = [
        value for value in append_values if value is not None and value > 0
    ]
    overlap_ratios = [
        feature.overlap_ratio_with_previous
        for feature in turn2plus
        if feature.overlap_ratio_with_previous is not None
    ]
    gaps = [
        feature.inter_turn_gap_s
        for feature in turn2plus
        if feature.inter_turn_gap_s is not None
    ]
    return {
        "request_count": len(requests),
        "session_count": len(session_turns),
        "multi_turn_session_count": sum(1 for turns in session_turns.values() if turns > 1),
        "turn2plus_count": len(turn2plus),
        "direct_eligible_turn2plus_count": len(direct_eligible),
        "direct_eligible_turn2plus_ratio": (
            len(direct_eligible) / len(turn2plus) if turn2plus else 0.0
        ),
        "turn_count_distribution": dict(sorted(Counter(session_turns.values()).items())),
        "request_type_distribution": dict(
            sorted(Counter(request.request_type for request in requests).items())
        ),
        "turn_id_distribution": dict(
            sorted(Counter(request.turn_id for request in requests).items())
        ),
        "append_bin_distribution": dict(
            sorted(Counter(feature.append_bin for feature in turn2plus).items())
        ),
        "input_bin_distribution": dict(
            sorted(Counter(feature.input_bin for feature in features.values()).items())
        ),
        "output_bin_distribution": dict(
            sorted(Counter(feature.output_bin for feature in features.values()).items())
        ),
        "resident_bin_distribution": dict(
            sorted(Counter(feature.resident_bin for feature in features.values()).items())
        ),
        "input_tokens_stats": _stats(
            [float(request.input_length) for request in requests]
        ),
        "output_tokens_stats": _stats(
            [float(request.output_length) for request in requests]
        ),
        "resident_tokens_stats": _stats(
            [float(feature.resident_tokens) for feature in features.values()]
        ),
        "append_input_tokens_stats": _stats(
            [float(value) for value in append_values if value is not None]
        ),
        "positive_append_input_tokens_stats": _stats(
            [float(value) for value in positive_append_values]
        ),
        "inter_turn_gap_s_stats": _stats([float(value) for value in gaps]),
        "overlap_ratio_stats": _stats([float(value) for value in overlap_ratios]),
        "non_positive_append_count": sum(
            1 for value in append_values if value is not None and value <= 0
        ),
    }


def _metrics_profile(
    rows: list[dict[str, Any]],
    features: dict[str, _TraceFeature],
) -> dict[str, Any]:
    return {
        "request_count": len(rows),
        "mechanism_distribution": dict(
            sorted(Counter(str(row.get("mechanism_name")) for row in rows).items())
        ),
        "execution_mode_distribution": dict(
            sorted(Counter(str(row.get("execution_mode")) for row in rows).items())
        ),
        "latency_s_stats": _stats(_numeric_values(rows, "latency_s")),
        "ttft_s_stats": _stats(_numeric_values(rows, "ttft_s")),
        "tpot_s_stats": _stats(_numeric_values(rows, "tpot_s")),
        "cached_tokens_stats": _stats(_numeric_values(rows, "cached_tokens")),
        "actual_kv_transfer_blocks_stats": _stats(
            _numeric_values(rows, "actual_kv_transfer_blocks")
        ),
        "session_reused_count": sum(1 for row in rows if row.get("session_reused")),
        "session_reset_count": sum(1 for row in rows if row.get("session_reset")),
        "error_count": sum(1 for row in rows if row.get("error") is not None),
        "by_turn_class": _group_metrics(rows, features, lambda feature, _row: feature.turn_class),
        "by_direct_eligible": _group_metrics(
            rows,
            features,
            lambda feature, _row: "eligible" if feature.direct_eligible else "not_eligible",
        ),
        "by_append_bin": _group_metrics(rows, features, lambda feature, _row: feature.append_bin),
        "by_resident_bin": _group_metrics(
            rows,
            features,
            lambda feature, _row: feature.resident_bin,
        ),
        "by_execution_mode": _group_metrics(
            rows,
            features,
            lambda _feature, row: str(row.get("execution_mode")),
        ),
    }


def _paired_comparison(
    *,
    baseline: list[dict[str, Any]],
    candidate: list[dict[str, Any]],
    features: dict[str, _TraceFeature],
) -> dict[str, Any]:
    baseline_by_id = {
        str(row.get("request_id")): row
        for row in baseline
        if row.get("latency_s") is not None
    }
    candidate_by_id = {
        str(row.get("request_id")): row
        for row in candidate
        if row.get("latency_s") is not None
    }
    paired_ids = sorted(set(baseline_by_id) & set(candidate_by_id))
    pairs = [
        (baseline_by_id[request_id], candidate_by_id[request_id], features.get(request_id))
        for request_id in paired_ids
    ]
    pairs = [pair for pair in pairs if pair[2] is not None]

    return {
        "overall": _delta_summary(pairs),
        "by_turn_class": _group_deltas(
            pairs,
            lambda feature, _base, _cand: feature.turn_class,
        ),
        "by_direct_eligible": _group_deltas(
            pairs,
            lambda feature, _base, _cand: (
                "eligible" if feature.direct_eligible else "not_eligible"
            ),
        ),
        "by_append_bin": _group_deltas(
            pairs,
            lambda feature, _base, _cand: feature.append_bin,
        ),
        "by_resident_bin": _group_deltas(
            pairs,
            lambda feature, _base, _cand: feature.resident_bin,
        ),
        "by_candidate_execution_mode": _group_deltas(
            pairs,
            lambda _feature, _base, cand: str(cand.get("execution_mode")),
        ),
    }


def _group_metrics(
    rows: list[dict[str, Any]],
    features: dict[str, _TraceFeature],
    key_fn,
) -> dict[str, Any]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        feature = features.get(str(row.get("request_id")))
        if feature is None:
            continue
        grouped[str(key_fn(feature, row))].append(row)
    return {
        key: {
            "count": len(group_rows),
            "latency_s_stats": _stats(_numeric_values(group_rows, "latency_s")),
            "ttft_s_stats": _stats(_numeric_values(group_rows, "ttft_s")),
            "tpot_s_stats": _stats(_numeric_values(group_rows, "tpot_s")),
            "session_reused_count": sum(
                1 for row in group_rows if row.get("session_reused")
            ),
            "error_count": sum(1 for row in group_rows if row.get("error") is not None),
        }
        for key, group_rows in sorted(grouped.items())
    }


def _group_deltas(
    pairs: list[tuple[dict[str, Any], dict[str, Any], _TraceFeature | None]],
    key_fn,
) -> dict[str, Any]:
    grouped: dict[str, list[tuple[dict[str, Any], dict[str, Any], _TraceFeature | None]]] = (
        defaultdict(list)
    )
    for base, cand, feature in pairs:
        if feature is None:
            continue
        grouped[str(key_fn(feature, base, cand))].append((base, cand, feature))
    return {key: _delta_summary(group_pairs) for key, group_pairs in sorted(grouped.items())}


def _delta_summary(
    pairs: list[tuple[dict[str, Any], dict[str, Any], _TraceFeature | None]],
) -> dict[str, Any]:
    latency_deltas = [
        float(cand["latency_s"]) - float(base["latency_s"])
        for base, cand, _feature in pairs
        if base.get("latency_s") is not None and cand.get("latency_s") is not None
    ]
    ttft_deltas = [
        float(cand["ttft_s"]) - float(base["ttft_s"])
        for base, cand, _feature in pairs
        if base.get("ttft_s") is not None and cand.get("ttft_s") is not None
    ]
    return {
        "paired_count": len(latency_deltas),
        "candidate_faster_count": sum(1 for delta in latency_deltas if delta < 0),
        "candidate_slower_count": sum(1 for delta in latency_deltas if delta > 0),
        "latency_delta_s_stats": _stats(latency_deltas),
        "ttft_delta_s_stats": _stats(ttft_deltas),
        "total_latency_delta_s": sum(latency_deltas),
        "mean_baseline_latency_s": _mean(
            [
                float(base["latency_s"])
                for base, cand, _feature in pairs
                if base.get("latency_s") is not None and cand.get("latency_s") is not None
            ]
        ),
        "mean_candidate_latency_s": _mean(
            [
                float(cand["latency_s"])
                for base, cand, _feature in pairs
                if base.get("latency_s") is not None and cand.get("latency_s") is not None
            ]
        ),
    }


def _load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                rows.append(json.loads(line))
    return rows


def _numeric_values(rows: Iterable[dict[str, Any]], key: str) -> list[float]:
    values: list[float] = []
    for row in rows:
        value = row.get(key)
        if value is not None:
            values.append(float(value))
    return values


def _stats(values: list[float]) -> dict[str, float] | None:
    clean = [value for value in values if value is not None]
    if not clean:
        return None
    clean.sort()
    return {
        "count": float(len(clean)),
        "mean": statistics.fmean(clean),
        "p50": _percentile(clean, 0.50),
        "p90": _percentile(clean, 0.90),
        "p99": _percentile(clean, 0.99),
        "min": clean[0],
        "max": clean[-1],
    }


def _percentile(sorted_values: list[float], percentile: float) -> float:
    if len(sorted_values) == 1:
        return sorted_values[0]
    index = round((len(sorted_values) - 1) * percentile)
    return sorted_values[index]


def _mean(values: list[float]) -> float | None:
    if not values:
        return None
    return statistics.fmean(values)


def _token_bin(value: int | None) -> str:
    if value is None:
        return "none"
    if value <= 0:
        return "<=0"
    if value <= 512:
        return "1-512"
    if value <= 2048:
        return "513-2048"
    if value <= 8192:
        return "2049-8192"
    if value <= 32768:
        return "8193-32768"
    return ">32768"