#!/usr/bin/env python3
"""Audit long code traces before choosing a replay window and max model length."""

from __future__ import annotations

import argparse
import json
import math
import statistics
from collections import Counter
from pathlib import Path
from typing import Any, Iterable, Sequence


BLOCK_SIZE_CANDIDATES = (16, 32, 64, 128, 256, 512, 1024)
MAX_MODEL_LEN_CANDIDATES = (40960, 65536, 98304, 131072, 262144)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--trace-root", type=Path)
    parser.add_argument("--source", type=Path, action="append")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--min-minutes", type=int, default=60)
    parser.add_argument("--max-minutes", type=int, default=75)
    parser.add_argument("--bin-seconds", type=int, default=60)
    parser.add_argument("--max-acceptable-gap-s", type=float, default=5.0)
    parser.add_argument("--model-position-limit", type=int, default=262144)
    return parser.parse_args()


def percentile(values: Sequence[int | float], fraction: float) -> float | None:
    if not values:
        return None
    ordered = sorted(float(value) for value in values)
    position = (len(ordered) - 1) * fraction
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    return ordered[lower] * (upper - position) + ordered[upper] * (position - lower)


def distribution(values: Sequence[int | float]) -> dict[str, int | float | None]:
    return {
        "count": len(values),
        "min": min(values) if values else None,
        "p50": percentile(values, 0.50),
        "p90": percentile(values, 0.90),
        "p95": percentile(values, 0.95),
        "p99": percentile(values, 0.99),
        "max": max(values) if values else None,
        "mean": statistics.fmean(values) if values else None,
    }


def parse_hash_ids(value: Any) -> list[Any]:
    if isinstance(value, list):
        return value
    if isinstance(value, str):
        stripped = value.strip()
        if not stripped:
            return []
        if stripped.startswith("["):
            decoded = json.loads(stripped)
            if not isinstance(decoded, list):
                raise ValueError("hash_ids JSON must decode to a list")
            return decoded
        delimiter = "|" if "|" in stripped else ","
        return [part for part in stripped.split(delimiter) if part.strip()]
    if value is None:
        return []
    return [value]


def iter_jsonl(path: Path) -> Iterable[tuple[int, dict[str, Any]]]:
    with path.open() as stream:
        for line_number, line in enumerate(stream, 1):
            if not line.strip():
                continue
            row = json.loads(line)
            if not isinstance(row, dict):
                raise ValueError(f"{path}:{line_number}: row must be an object")
            yield line_number, row


def choose_window(
    *,
    counts: Sequence[int],
    max_gaps: Sequence[float],
    first_timestamp: float,
    min_minutes: int,
    max_minutes: int,
    bin_seconds: int,
    max_acceptable_gap_s: float,
) -> dict[str, Any] | None:
    candidates = []
    for minutes in range(max_minutes, min_minutes - 1, -1):
        bins = math.ceil(minutes * 60 / bin_seconds)
        for start_bin in range(0, len(counts) - bins + 1):
            selected = counts[start_bin : start_bin + bins]
            mean = statistics.fmean(selected)
            cv = statistics.pstdev(selected) / mean if mean else math.inf
            max_gap = max(max_gaps[start_bin : start_bin + bins], default=0.0)
            candidates.append(
                {
                    "_score": (
                        max_gap > max_acceptable_gap_s,
                        cv,
                        max_gap,
                        -minutes,
                        start_bin,
                    ),
                    "start_bin": start_bin,
                    "minutes": minutes,
                    "count_mean_per_bin": mean,
                    "count_cv": cv,
                    "count_min_per_bin": min(selected),
                    "count_max_per_bin": max(selected),
                    "max_gap_s": max_gap,
                }
            )
    if not candidates:
        return None
    chosen = min(candidates, key=lambda item: item["_score"])
    chosen.pop("_score")
    chosen["start_timestamp"] = first_timestamp + chosen["start_bin"] * bin_seconds
    chosen["end_timestamp"] = chosen["start_timestamp"] + chosen["minutes"] * 60
    return chosen


def scan_source(path: Path, args: argparse.Namespace) -> dict[str, Any]:
    rows = 0
    source_rows = 0
    invalid_zero_token_rows = 0
    invalid_zero_token_examples: list[dict[str, Any]] = []
    first_timestamp = None
    last_timestamp = None
    previous_timestamp = None
    counts: Counter[int] = Counter()
    max_gaps: dict[int, float] = {}
    input_lengths: list[int] = []
    output_lengths: list[int] = []
    total_lengths: list[int] = []
    hash_rows = 0
    hash_matches = Counter()
    prompt_rows = 0
    sampling_rows = 0
    schema_keys: Counter[str] = Counter()
    for line_number, row in iter_jsonl(path):
        source_rows += 1
        missing = [
            key
            for key in ("timestamp", "input_length", "output_length")
            if key not in row
        ]
        if missing:
            raise ValueError(f"{path}:{line_number}: missing required fields {missing}")
        timestamp = float(row["timestamp"])
        input_tokens = int(row["input_length"])
        output_tokens = int(row["output_length"])
        schema_keys.update(row.keys())
        if input_tokens <= 0 or output_tokens <= 0:
            invalid_zero_token_rows += 1
            if len(invalid_zero_token_examples) < 20:
                invalid_zero_token_examples.append(
                    {
                        "line_number": line_number,
                        "chat_id": row.get("chat_id"),
                        "timestamp": timestamp,
                        "input_length": input_tokens,
                        "output_length": output_tokens,
                    }
                )
            continue
        if first_timestamp is None:
            first_timestamp = timestamp
        if previous_timestamp is not None and timestamp < previous_timestamp:
            raise ValueError(
                f"{path}:{line_number}: timestamp {timestamp} < {previous_timestamp}"
            )
        bin_index = math.floor((timestamp - first_timestamp) / args.bin_seconds)
        counts[bin_index] += 1
        if previous_timestamp is not None:
            previous_bin = math.floor(
                (previous_timestamp - first_timestamp) / args.bin_seconds
            )
            max_gaps[previous_bin] = max(
                max_gaps.get(previous_bin, 0.0),
                timestamp - previous_timestamp,
            )
        input_lengths.append(input_tokens)
        output_lengths.append(output_tokens)
        total_lengths.append(input_tokens + output_tokens)
        hashes = parse_hash_ids(row.get("hash_ids"))
        if hashes:
            hash_rows += 1
            for block_size in BLOCK_SIZE_CANDIDATES:
                if len(hashes) == math.ceil(input_tokens / block_size):
                    hash_matches[block_size] += 1
        prompt_rows += int(
            isinstance(row.get("prompt"), (str, list)) and bool(row.get("prompt"))
        )
        sampling_rows += int("sampling_u" in row)
        rows += 1
        previous_timestamp = timestamp
        last_timestamp = timestamp
    if not rows or first_timestamp is None or last_timestamp is None:
        raise ValueError(f"{path}: empty trace")
    total_bins = math.floor((last_timestamp - first_timestamp) / args.bin_seconds) + 1
    chosen = choose_window(
        counts=[counts[index] for index in range(total_bins)],
        max_gaps=[max_gaps.get(index, 0.0) for index in range(total_bins)],
        first_timestamp=first_timestamp,
        min_minutes=args.min_minutes,
        max_minutes=args.max_minutes,
        bin_seconds=args.bin_seconds,
        max_acceptable_gap_s=args.max_acceptable_gap_s,
    )
    return {
        "source": str(path.resolve()),
        "rows": rows,
        "source_rows": source_rows,
        "invalid_zero_token_rows": invalid_zero_token_rows,
        "invalid_zero_token_fraction": invalid_zero_token_rows / source_rows,
        "invalid_zero_token_examples": invalid_zero_token_examples,
        "first_timestamp": first_timestamp,
        "last_timestamp": last_timestamp,
        "span_s": last_timestamp - first_timestamp,
        "request_rate_per_s": rows / max(last_timestamp - first_timestamp, 1.0),
        "input_length": distribution(input_lengths),
        "output_length": distribution(output_lengths),
        "total_length": distribution(total_lengths),
        "over_max_model_len": {
            str(limit): {
                "requests": sum(value > limit for value in total_lengths),
                "fraction": sum(value > limit for value in total_lengths) / rows,
            }
            for limit in MAX_MODEL_LEN_CANDIDATES
        },
        "hash_contract": {
            "rows_with_hash_ids": hash_rows,
            "candidate_exact_match_rows": {
                str(size): hash_matches[size] for size in BLOCK_SIZE_CANDIDATES
            },
            "exact_source_block_size": next(
                (
                    size
                    for size in BLOCK_SIZE_CANDIDATES
                    if hash_rows and hash_matches[size] == hash_rows
                ),
                None,
            ),
        },
        "prompt_rows": prompt_rows,
        "sampling_u_rows": sampling_rows,
        "schema_field_counts": dict(sorted(schema_keys.items())),
        "stable_window": chosen,
    }


def scan_window(source: Path, window: dict[str, Any]) -> dict[str, Any]:
    start = float(window["start_timestamp"])
    end = float(window["end_timestamp"])
    inputs: list[int] = []
    outputs: list[int] = []
    totals: list[int] = []
    for _, row in iter_jsonl(source):
        timestamp = float(row["timestamp"])
        if timestamp < start:
            continue
        if timestamp >= end:
            break
        input_tokens = int(row["input_length"])
        output_tokens = int(row["output_length"])
        if input_tokens <= 0 or output_tokens <= 0:
            continue
        inputs.append(input_tokens)
        outputs.append(output_tokens)
        totals.append(input_tokens + output_tokens)
    return {
        "requests": len(totals),
        "input_length": distribution(inputs),
        "output_length": distribution(outputs),
        "total_length": distribution(totals),
        "max_model_len_coverage": {
            str(limit): {
                "covered_requests": sum(value <= limit for value in totals),
                "excluded_requests": sum(value > limit for value in totals),
                "coverage": sum(value <= limit for value in totals) / len(totals),
            }
            for limit in MAX_MODEL_LEN_CANDIDATES
        },
    }


def resolve_sources(args: argparse.Namespace) -> list[Path]:
    if args.source:
        return [path.resolve() for path in args.source]
    if args.trace_root is None:
        raise ValueError("provide --trace-root or one or more --source")
    sources = sorted(
        path.resolve()
        for path in args.trace_root.glob("*.jsonl")
        if "prompt" not in path.stem.lower()
    )
    if not sources:
        raise FileNotFoundError(f"no non-prompt JSONL files under {args.trace_root}")
    return sources


def main() -> None:
    args = parse_args()
    if not 0 < args.min_minutes <= args.max_minutes:
        raise ValueError("require 0 < min_minutes <= max_minutes")
    sources = resolve_sources(args)
    files = [scan_source(path, args) for path in sources]
    eligible = [item for item in files if item["stable_window"] is not None]
    if not eligible:
        chosen = None
        data_gate = "BLOCKED_NO_1H_WINDOW"
    else:
        chosen = min(
            eligible,
            key=lambda item: (
                item["stable_window"]["max_gap_s"] > args.max_acceptable_gap_s,
                item["stable_window"]["count_cv"],
                -item["stable_window"]["minutes"],
                item["source"],
            ),
        )
        chosen["selected_window_stats"] = scan_window(
            Path(chosen["source"]), chosen["stable_window"]
        )
        exact_block_size = chosen["hash_contract"]["exact_source_block_size"]
        max_total = chosen["selected_window_stats"]["total_length"]["max"]
        data_gate = (
            "PASS"
            if exact_block_size is not None
            and max_total is not None
            and max_total <= args.model_position_limit
            else "BLOCKED_HASH_OR_POSITION_CONTRACT"
        )
    recommendation = None
    if chosen is not None:
        maximum = chosen["selected_window_stats"]["total_length"]["max"]
        recommendation = next(
            (
                limit
                for limit in MAX_MODEL_LEN_CANDIDATES
                if maximum <= limit <= args.model_position_limit
            ),
            None,
        )
    payload = {
        "schema": "frontier-code-trace-audit-v1",
        "trace_root": str(args.trace_root.resolve()) if args.trace_root else None,
        "sources": [str(path) for path in sources],
        "window_policy": {
            "min_minutes": args.min_minutes,
            "max_minutes": args.max_minutes,
            "bin_seconds": args.bin_seconds,
            "max_acceptable_gap_s": args.max_acceptable_gap_s,
            "selection": "lowest density CV after rejecting anomalous-gap windows",
        },
        "model_position_limit": args.model_position_limit,
        "files": files,
        "selected": chosen,
        "max_model_len_recommendation": recommendation,
        "data_gate": data_gate,
        "runtime_gate": (
            "PENDING: vLLM startup must prove enough KV blocks and nonzero "
            "max concurrency at the recommended max_model_len for each TP"
        ),
    }
    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({"data_gate": data_gate, "output": str(args.output)}))


if __name__ == "__main__":
    main()