#!/usr/bin/env python3
"""Freeze the exact request cohort and its operator-profile support.

This calls AITuner's production trace loader, including its input-length
filter, uniform max-request downsampling, output override, and sampling-u
threshold semantics.  Prefix reuse below is a no-eviction upper bound; the
actual cache state remains scheduler/config dependent.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
from pathlib import Path
from typing import Any, Iterable

from aituner.spec import load_study_spec
from aituner.trace import load_trace_requests, select_requests_for_threshold


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--study", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--cohort-output", type=Path, required=True)
    parser.add_argument(
        "--thresholds",
        type=float,
        nargs="+",
        default=[0.125, 0.25, 0.5, 1.0],
    )
    return parser.parse_args()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def nearest_rank(values: Iterable[float], percentile: float) -> float:
    ordered = sorted(values)
    if not ordered:
        return 0.0
    index = min(
        len(ordered) - 1,
        max(0, math.ceil(percentile / 100.0 * len(ordered)) - 1),
    )
    return float(ordered[index])


def distribution(values: list[float]) -> dict[str, float]:
    return {
        "min": min(values, default=0.0),
        "p10": nearest_rank(values, 10),
        "p25": nearest_rank(values, 25),
        "p50": nearest_rank(values, 50),
        "p75": nearest_rank(values, 75),
        "p90": nearest_rank(values, 90),
        "p95": nearest_rank(values, 95),
        "p99": nearest_rank(values, 99),
        "max": max(values, default=0.0),
    }


def prefix_upper_bound(requests: list[Any], block_size: int) -> dict[str, Any]:
    seen: set[Any] = set()
    total_blocks = 0
    reusable_blocks = 0
    leading_reusable_blocks: list[float] = []
    reusable_tokens: list[float] = []
    rows_with_hashes = 0
    for request in requests:
        hashes = request.metadata.get("hash_ids")
        if not isinstance(hashes, list):
            continue
        rows_with_hashes += 1
        total_blocks += len(hashes)
        reusable = sum(hash_id in seen for hash_id in hashes)
        reusable_blocks += reusable
        leading = 0
        for hash_id in hashes:
            if hash_id not in seen:
                break
            leading += 1
        leading_reusable_blocks.append(float(leading))
        reusable_tokens.append(
            float(min(request.prompt_tokens_hint or 0, leading * block_size))
        )
        seen.update(hashes)
    return {
        "semantics": (
            "arrival-ordered, infinite-capacity/no-eviction upper bound; "
            "not an observed KV-cache hit rate"
        ),
        "rows_with_hash_ids": rows_with_hashes,
        "total_blocks": total_blocks,
        "unique_blocks": len(seen),
        "any_position_reusable_block_ratio": (
            reusable_blocks / total_blocks if total_blocks else 0.0
        ),
        "leading_reusable_blocks_per_request": distribution(
            leading_reusable_blocks
        ),
        "leading_reusable_tokens_per_request": distribution(reusable_tokens),
    }


def summarize(requests: list[Any], block_size: int) -> dict[str, Any]:
    input_lengths = [float(request.prompt_tokens_hint or 0) for request in requests]
    output_lengths = [
        float(request.completion_tokens_hint or 0) for request in requests
    ]
    hash_counts = [
        float(len(request.metadata.get("hash_ids") or [])) for request in requests
    ]
    arrivals = [request.arrival_s for request in requests]
    interarrivals = [
        max(0.0, arrivals[index] - arrivals[index - 1])
        for index in range(1, len(arrivals))
    ]
    return {
        "request_count": len(requests),
        "input_tokens": distribution(input_lengths),
        "output_tokens": distribution(output_lengths),
        "hash_blocks_per_request": distribution(hash_counts),
        "interarrival_s": distribution(interarrivals),
        "sampling_u": distribution([request.sampling_u for request in requests]),
        "multi_turn_fraction": (
            sum(
                isinstance(request.metadata.get("turn"), (int, float))
                and request.metadata["turn"] > 1
                for request in requests
            )
            / len(requests)
            if requests
            else 0.0
        ),
        "prefix_reuse_upper_bound": prefix_upper_bound(requests, block_size),
    }


def cohort_row(request: Any) -> dict[str, Any]:
    return {
        "row_id": request.row_id,
        "arrival_s": request.arrival_s,
        "sampling_u": request.sampling_u,
        "input_tokens": request.prompt_tokens_hint,
        "output_tokens": request.completion_tokens_hint,
        "hash_ids": request.metadata.get("hash_ids"),
        "turn": request.metadata.get("turn"),
        "parent_chat_id": request.metadata.get("parent_chat_id"),
        "type": request.metadata.get("type"),
    }


def main() -> None:
    args = parse_args()
    study = load_study_spec(args.study)
    window, cohort = load_trace_requests(study, study_spec_path=args.study)
    block_size = int(window.source_payload.get("block_size") or 1)

    args.cohort_output.parent.mkdir(parents=True, exist_ok=True)
    with args.cohort_output.open("w", encoding="utf-8") as handle:
        for request in cohort:
            handle.write(json.dumps(cohort_row(request), sort_keys=True) + "\n")

    payload = {
        "schema_version": "qwen30_trace_profile_support.v1",
        "study": str(args.study.resolve()),
        "study_sha256": sha256_file(args.study),
        "trace": str(window.trace_path),
        "trace_sha256": sha256_file(window.trace_path),
        "window_id": window.window_id,
        "block_size": block_size,
        "loader_contract": {
            "input_length_filter": {
                "min": study.trace.input_length_filter.min_input_tokens,
                "max": study.trace.input_length_filter.max_input_tokens,
            }
            if study.trace.input_length_filter is not None
            else None,
            "completion_tokens_override": study.trace.completion_tokens_override,
            "max_requests_per_probe": study.trace.max_requests_per_probe,
            "replay_time_scale": study.trace.replay_time_scale,
            "ordering": "arrival_s",
            "downsampling": "AITuner _downsample_requests before sampling_u threshold",
        },
        "full_downsampled_cohort": summarize(cohort, block_size),
        "threshold_cohorts": {
            str(threshold): summarize(
                select_requests_for_threshold(cohort, threshold=threshold),
                block_size,
            )
            for threshold in args.thresholds
        },
        "limits": [
            "Trace length/hash support does not determine dynamic decode or mixed batch shapes.",
            "Those shapes depend jointly on arrival history, SLO pressure, TP execution time, MNS, chunking, and KV eviction.",
            "MoE expert routing is not present in the trace and must be measured from model execution.",
        ],
    }
    payload["cohort_output"] = str(args.cohort_output.resolve())
    payload["cohort_sha256"] = sha256_file(args.cohort_output)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps(payload, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()
