#!/usr/bin/env python3
"""Phase-5 private-manifest transforms and timestamp-scheduled P3 client wrapper."""

from __future__ import annotations

import argparse
import asyncio
import json
import math
import sys
from pathlib import Path
from typing import Any

import aiohttp

import opprof_phase3_client as p3


def _numeric(values: list[float | int]) -> dict[str, Any]:
    finite = [float(value) for value in values if math.isfinite(float(value))]
    return {
        "n": len(values),
        "finite_n": len(finite),
        "missing_n": len(values) - len(finite),
        "min": min(finite) if finite else None,
        "max": max(finite) if finite else None,
        "distinct_n": len(set(finite)),
        "sum": sum(finite),
    }


def _r16(rows: list[dict[str, Any]]) -> float:
    groups = [rows[index : index + 16] for index in range(0, len(rows) - 15, 16)]
    useful = sum(sum(int(row["input_tokens"]) for row in group) for group in groups)
    rectangular = sum(16 * max(int(row["input_tokens"]) for row in group) for group in groups)
    return 1.0 - useful / rectangular


def _source_timestamps(path: Path, indices: set[int], field: str) -> dict[int, float]:
    result: dict[int, float] = {}
    maximum = max(indices)
    with path.open(encoding="utf-8") as source:
        for index, line in enumerate(source):
            if index in indices:
                value = float(json.loads(line)[field])
                if not math.isfinite(value):
                    raise ValueError(f"non-finite source timestamp at {index}")
                result[index] = value
            if index >= maximum:
                break
    if set(result) != indices:
        raise ValueError("timestamp source did not cover all source_index values")
    return result


def transform(args: argparse.Namespace) -> dict[str, Any]:
    rows = p3.load_manifest(Path(args.input))[: args.take_first]
    if len(rows) != args.take_first:
        raise ValueError(f"requested {args.take_first} rows, found {len(rows)}")
    indices = [int(row[args.join_key]) for row in rows]
    timestamps = _source_timestamps(
        Path(args.timestamp_source), set(indices), args.timestamp_field
    )
    source_times = [timestamps[index] for index in indices]
    if any(right < left for left, right in zip(source_times, source_times[1:])):
        raise ValueError("selected timestamps are not nondecreasing")
    if source_times[-1] <= source_times[0]:
        raise ValueError("selected timestamp span is not positive")
    end_s = (len(rows) - 1) / args.target_rate
    scale = end_s / (source_times[-1] - source_times[0])
    recorded_slots = [(value - source_times[0]) * scale for value in source_times]
    uniform_slots = [index / args.target_rate for index in range(len(rows))]
    slots = recorded_slots if args.arrival == "recorded-scaled" else uniform_slots

    for index, row in enumerate(rows):
        row["arrival"] = args.arrival
        row["arrival_s"] = slots[index]
        row["original_index"] = index
        row["source_timestamp"] = source_times[index]

    original = list(rows)
    max_added_delay = 0.0
    if args.service_order == "length-binned":
        edges = [int(item) for item in args.length_bin_edges.split(",")]

        def bin_id(row: dict[str, Any]) -> int:
            length = int(row["input_tokens"])
            for index, edge in enumerate(edges):
                if length <= edge:
                    return index
            raise ValueError(f"input length {length} exceeds final edge")

        reordered: list[dict[str, Any]] = []
        for offset in range(0, len(rows), args.reorder_block_size):
            block = rows[offset : offset + args.reorder_block_size]
            ordered = sorted(
                block,
                key=lambda row: (
                    bin_id(row),
                    int(row["input_tokens"]),
                    int(row["original_index"]),
                ),
            )
            block_slots = slots[offset : offset + len(block)]
            for position, row in enumerate(ordered):
                added = max(0.0, block_slots[position] - float(row["arrival_s"]))
                max_added_delay = max(max_added_delay, added)
                row["arrival_s"] = block_slots[position]
            reordered.extend(ordered)
        rows = reordered
        if max_added_delay > args.max_added_delay_seconds + 1e-9:
            raise ValueError(
                f"fairness cap exceeded: {max_added_delay} > {args.max_added_delay_seconds}"
            )
    elif args.service_order != "original":
        raise ValueError(f"unsupported service order: {args.service_order}")

    if sorted(row["request_id"] for row in rows) != sorted(
        row["request_id"] for row in original
    ):
        raise AssertionError("request identity changed")
    for key in ("input_tokens", "output_tokens"):
        if sum(int(row[key]) for row in rows) != sum(int(row[key]) for row in original):
            raise AssertionError(f"{key} total changed")
    arrival_values = [float(row["arrival_s"]) for row in rows]
    if any(right < left for left, right in zip(arrival_values, arrival_values[1:])):
        raise AssertionError("assigned arrival slots are not nondecreasing")

    output = Path(args.out)
    p3.atomic_jsonl(output, rows, mode=0o600)
    summary = {
        "schema": 1,
        "path": str(output),
        "sha256": p3.sha256_file(output),
        "rows": len(rows),
        "arrival": args.arrival,
        "service_order": args.service_order,
        "target_rate": args.target_rate,
        "input_tokens": _numeric([int(row["input_tokens"]) for row in rows]),
        "output_tokens": _numeric([int(row["output_tokens"]) for row in rows]),
        "arrival_s": _numeric(arrival_values),
        "source_timestamp": _numeric(source_times),
        "r16": _r16(rows),
        "max_added_delay_seconds": max_added_delay,
        "request_id_set_sha256": p3.hashlib.sha256(
            "\n".join(sorted(str(row["request_id"]) for row in rows)).encode()
        ).hexdigest(),
        "invariants": {
            "same_request_ids": True,
            "same_input_tokens": True,
            "same_output_tokens": True,
            "arrival_nondecreasing": True,
            "fairness_cap": max_added_delay <= args.max_added_delay_seconds + 1e-9,
            "no_prompt_in_summary": True,
        },
    }
    p3.atomic_json(output.with_suffix(output.suffix + ".summary.json"), summary, mode=0o600)
    print(json.dumps(summary, sort_keys=True))
    return summary


async def finite_timestamp_load(
    ctx: p3.RunContext, session: aiohttp.ClientSession, rate: float
) -> list[dict[str, Any]]:
    if "arrival_s" not in ctx.rows[0]:
        return await _ORIGINAL_FINITE_LOAD(ctx, session, rate)
    sem = asyncio.Semaphore(ctx.args.max_concurrency)
    tasks: list[asyncio.Task[dict[str, Any]]] = []

    async def limited(row: dict[str, Any], scheduled: float) -> dict[str, Any]:
        async with sem:
            return await p3.request_one(ctx, session, row, scheduled)

    for expected in ctx.rows:
        scheduled = ctx.t0 + float(expected["arrival_s"])
        delay = scheduled - asyncio.get_running_loop().time()
        if delay > 0:
            try:
                await asyncio.wait_for(ctx.stop_event.wait(), timeout=delay)
                break
            except asyncio.TimeoutError:
                pass
        if ctx.stop_event.is_set():
            break
        row = await ctx.next_row()
        if row["request_id"] != expected["request_id"]:
            raise AssertionError("timestamp scheduler row drift")
        tasks.append(asyncio.create_task(limited(row, scheduled)))
    return await asyncio.gather(*tasks) if tasks else []


_ORIGINAL_FINITE_LOAD = p3.finite_load
p3.finite_load = finite_timestamp_load


def build_transform_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser.add_argument("--in", dest="input", required=True)
    parser.add_argument("--take-first", type=int, required=True)
    parser.add_argument("--timestamp-source", required=True)
    parser.add_argument("--join-key", default="source_index")
    parser.add_argument("--timestamp-field", default="timestamp")
    parser.add_argument("--arrival", choices=("recorded-scaled", "uniform"), required=True)
    parser.add_argument("--target-rate", type=float, required=True)
    parser.add_argument("--service-order", choices=("original", "length-binned"), required=True)
    parser.add_argument("--reorder-block-size", type=int, default=32)
    parser.add_argument("--analysis-cohort-size", type=int, default=16)
    parser.add_argument("--length-bin-edges", default="512,1024,2048,4096,8192,16384,32768")
    parser.add_argument("--max-added-delay-seconds", type=float, default=64)
    parser.add_argument("--out", required=True)
    return parser


def main() -> None:
    if len(sys.argv) > 1 and sys.argv[1] == "transform":
        transform(build_transform_parser().parse_args(sys.argv[2:]))
        return
    fixed_rate = None
    if "--fixed-request-rate" in sys.argv:
        index = sys.argv.index("--fixed-request-rate")
        fixed_rate = float(sys.argv[index + 1])
        del sys.argv[index : index + 2]
    args = p3.build_parser().parse_args()
    if args.command != "run":
        p3.main()
        return
    if fixed_rate is not None:
        if args.load_point != "moderate" or fixed_rate <= 0 or not math.isfinite(fixed_rate):
            raise ValueError("--fixed-request-rate requires positive finite moderate rate")
        result_dir = Path(args.result_dir)
        result_dir.mkdir(parents=True, exist_ok=True)
        source = result_dir / "fixed-rate-source.json"
        p3.atomic_json(source, {"clean": {"completed_throughput_rps": fixed_rate}})
        args.saturation_result = str(source)
        args.rate_fraction = 1.0
    if args.profile_after_clean and not args.profile_trace_dir:
        raise ValueError("--profile-after-clean requires --profile-trace-dir")
    print(json.dumps(asyncio.run(p3.run_load(args)), sort_keys=True))


if __name__ == "__main__":
    main()
