#!/usr/bin/env python3
"""Prepare exact Qwen trace replays with block-16 prefix identities.

Prompt text is written only below ``private/``. Public manifests and Frontier
fixtures contain lengths, arrivals, session IDs, and deterministic block IDs.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import platform
import socket
import time
from pathlib import Path
from typing import Any


CSV_FIELDS = (
    "arrived_at",
    "num_prefill_tokens",
    "num_decode_tokens",
    "session_id",
    "block_hash_ids",
)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def threshold_name(value: float) -> str:
    return f"u{value:.12g}".replace(".", "p")


def token_payload(tokens: list[int]) -> bytes:
    return len(tokens).to_bytes(2, "little") + b"".join(
        int(token).to_bytes(4, "little", signed=False) for token in tokens
    )


def block_identity_records(
    token_ids: list[int], block_size: int
) -> list[tuple[int, bytes]]:
    """Return parent-sensitive identities and independent collision witnesses."""
    parent = b"FRONTIER_EXACT_TRACE_ROOT"
    records = []
    for start in range(0, len(token_ids), block_size):
        payload = token_payload(token_ids[start : start + block_size])
        identity_input = parent + b"\0" + payload
        parent = hashlib.blake2b(identity_input, digest_size=16).digest()
        records.append(
            (
                int.from_bytes(parent, "big", signed=False),
                hashlib.sha256(identity_input).digest(),
            )
        )
    return records


def block_identities(token_ids: list[int], block_size: int) -> list[int]:
    """Return parent-sensitive identities with the same prefix equivalence as vLLM."""
    return [identity for identity, _ in block_identity_records(token_ids, block_size)]


def root_sessions(rows: list[dict[str, Any]]) -> dict[int, int]:
    roots: dict[int, int] = {}
    for row in rows:
        chat_id = int(row["chat_id"])
        parent = int(row["parent_chat_id"])
        roots[chat_id] = chat_id if parent == -1 else roots.get(parent, parent)
    return roots


def update_digest(digest: Any, values: list[Any]) -> None:
    digest.update(json.dumps(values, separators=(",", ":")).encode())
    digest.update(b"\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--trace", type=Path, required=True)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--served-model-name", default="qwen3-30b-a3b")
    parser.add_argument("--sampling-u-max", type=float, action="append", required=True)
    parser.add_argument("--max-model-len", type=int, default=40960)
    parser.add_argument("--source-block-size", type=int, default=64)
    parser.add_argument("--runtime-block-size", type=int, default=16)
    parser.add_argument("--batch-size", type=int, default=16)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    thresholds = sorted(set(args.sampling_u_max))
    if not thresholds or thresholds[0] < 0 or thresholds[-1] > 1:
        raise ValueError("sampling thresholds must be in [0, 1]")
    if args.source_block_size % args.runtime_block_size:
        raise ValueError("source block size must be divisible by runtime block size")
    if min(args.max_model_len, args.source_block_size, args.runtime_block_size, args.batch_size) <= 0:
        raise ValueError("length and batch arguments must be positive")

    import transformers
    from transformers import AutoTokenizer

    rows = [json.loads(line) for line in args.trace.open() if line.strip()]
    roots = root_sessions(rows)
    eligible = [
        (index, row)
        for index, row in enumerate(rows)
        if int(row["input_length"]) + int(row["output_length"]) <= args.max_model_len
        and int(row["output_length"]) > 0
    ]
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    prepared: list[dict[str, Any]] = []
    identity_to_digest: dict[int, bytes] = {}
    source_to_runtime: dict[int, tuple[int, ...]] = {}
    runtime_to_source: dict[tuple[int, ...], int] = {}
    identity_collision_count = 0
    source_to_runtime_conflicts = 0
    runtime_to_source_conflicts = 0
    length_mismatches = 0
    source_hash_count_mismatches = 0
    started = time.time()

    for start in range(0, len(eligible), args.batch_size):
        batch = eligible[start : start + args.batch_size]
        encoded = tokenizer(
            [row["prompt"] for _, row in batch],
            add_special_tokens=False,
            padding=False,
            truncation=False,
        )["input_ids"]
        for (source_index, row), token_ids in zip(batch, encoded, strict=True):
            token_ids = [int(token) for token in token_ids]
            if len(token_ids) != int(row["input_length"]):
                length_mismatches += 1
            source_hashes = [int(value) for value in row["hash_ids"]]
            if len(source_hashes) != math.ceil(len(token_ids) / args.source_block_size):
                source_hash_count_mismatches += 1
            identity_records = block_identity_records(token_ids, args.runtime_block_size)
            runtime_ids = [identity for identity, _ in identity_records]
            for runtime_id, witness in identity_records:
                previous = identity_to_digest.setdefault(runtime_id, witness)
                identity_collision_count += int(previous != witness)

            blocks_per_source = args.source_block_size // args.runtime_block_size
            for block_index, source_id in enumerate(source_hashes):
                begin = block_index * blocks_per_source
                relation = tuple(runtime_ids[begin : begin + blocks_per_source])
                previous_relation = source_to_runtime.setdefault(source_id, relation)
                source_to_runtime_conflicts += int(previous_relation != relation)
                previous_source = runtime_to_source.setdefault(relation, source_id)
                runtime_to_source_conflicts += int(previous_source != source_id)

            prepared.append(
                {
                    "source_index": source_index,
                    "chat_id": int(row["chat_id"]),
                    "session_id": roots[int(row["chat_id"])],
                    "arrival": float(row["timestamp"]),
                    "input_length": int(row["input_length"]),
                    "output_length": int(row["output_length"]),
                    "sampling_u": float(row["sampling_u"]),
                    "prompt": row["prompt"],
                    "runtime_block_ids": runtime_ids,
                }
            )

    args.output_root.mkdir(parents=True, exist_ok=True)
    threshold_manifests = []
    for threshold in thresholds:
        name = threshold_name(threshold)
        public = args.output_root / "public" / name
        private = args.output_root / "private" / name
        public.mkdir(parents=True, exist_ok=True)
        private.mkdir(parents=True, exist_ok=True)
        selected = [row for row in prepared if row["sampling_u"] <= threshold]
        frontier_path = public / "frontier.csv"
        real_path = private / "real_requests.jsonl"
        vector_digest = hashlib.sha256()
        with frontier_path.open("w", newline="") as output:
            writer = csv.DictWriter(output, fieldnames=CSV_FIELDS, lineterminator="\n")
            writer.writeheader()
            for row in selected:
                writer.writerow(
                    {
                        "arrived_at": f"{row['arrival']:.12f}",
                        "num_prefill_tokens": row["input_length"],
                        "num_decode_tokens": row["output_length"],
                        "session_id": row["session_id"],
                        "block_hash_ids": "|".join(
                            str(value) for value in row["runtime_block_ids"]
                        ),
                    }
                )
                update_digest(
                    vector_digest,
                    [
                        row["source_index"],
                        row["arrival"],
                        row["input_length"],
                        row["output_length"],
                        row["session_id"],
                        row["runtime_block_ids"],
                    ],
                )
        with real_path.open("w") as output:
            for row in selected:
                output.write(
                    json.dumps(
                        {
                            "source_index": row["source_index"],
                            "arrived_at": row["arrival"],
                            "input_length": row["input_length"],
                            "output_length": row["output_length"],
                            "session_id": row["session_id"],
                            "runtime_block_ids": row["runtime_block_ids"],
                            "body": {
                                "model": args.served_model_name,
                                "prompt": row["prompt"],
                                "min_tokens": row["output_length"],
                                "max_tokens": row["output_length"],
                                "ignore_eos": True,
                                "stream": True,
                            },
                        },
                        separators=(",", ":"),
                    )
                    + "\n"
                )
        threshold_manifests.append(
            {
                "sampling_u_max": threshold,
                "selected_requests": len(selected),
                "request_rate": len(selected) / 600.0,
                "frontier_csv": str(frontier_path.resolve()),
                "frontier_csv_sha256": sha256(frontier_path),
                "real_requests_private": str(real_path.resolve()),
                "real_requests_private_sha256": sha256(real_path),
                "row_vector_sha256": vector_digest.hexdigest(),
            }
        )

    failures = {
        "input_length_mismatches": length_mismatches,
        "source_hash_count_mismatches": source_hash_count_mismatches,
        "runtime_identity_collisions": identity_collision_count,
        "source_to_runtime_relation_conflicts": source_to_runtime_conflicts,
        "runtime_to_source_relation_conflicts": runtime_to_source_conflicts,
    }
    manifest = {
        "schema": "qwen30-exact-trace-block16-v1",
        "status": "pass" if not any(failures.values()) else "fail",
        "execution": {
            "host": socket.gethostname(),
            "python": platform.python_version(),
            "transformers": transformers.__version__,
            "tokenizer": type(tokenizer).__name__,
            "elapsed_seconds": round(time.time() - started, 3),
        },
        "source": {
            "trace": str(args.trace.resolve()),
            "trace_sha256": sha256(args.trace),
            "requests": len(rows),
            "eligible_requests": len(eligible),
            "model": str(args.model.resolve()),
            "max_model_len": args.max_model_len,
        },
        "block_contract": {
            "source_block_size": args.source_block_size,
            "runtime_block_size": args.runtime_block_size,
            "identity": "BLAKE2b-128(parent runtime identity, exact token-id block)",
            "unique_runtime_identities": len(identity_to_digest),
            "unique_source_relations": len(source_to_runtime),
            "failures": failures,
        },
        "selection_contract": (
            "eligible universe followed by source sampling_u threshold; no length "
            "selection or output override; original arrivals/order preserved"
        ),
        "thresholds": threshold_manifests,
        "privacy": "prompt text exists only under private/ and must not be harvested",
    }
    manifest_path = args.output_root / "public" / "manifest.json"
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"status": manifest["status"], "failures": failures, "thresholds": threshold_manifests}, sort_keys=True))
    if manifest["status"] != "pass":
        raise SystemExit(1)


if __name__ == "__main__":
    main()
