#!/usr/bin/env python3
"""Audit Qwen235B trace token lengths and source prefix-hash identities."""

from __future__ import annotations

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


def sha256_file(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 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", type=Path, 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("--batch-size", type=int, default=16)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.max_model_len <= 0 or args.source_block_size <= 0 or args.batch_size <= 0:
        raise ValueError("length and batch-size arguments must be positive")

    import transformers
    from transformers import AutoTokenizer

    rows = [json.loads(line) for line in args.trace.open() if line.strip()]
    context_exceeded = [
        row
        for row in rows
        if int(row["input_length"]) + int(row["output_length"]) > args.max_model_len
    ]
    zero_output = [row for row in rows if int(row["output_length"]) == 0]
    eligible = [
        row
        for row in 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)
    length_mismatch_count = 0
    hash_count_mismatch_count = 0
    hash_to_key: dict[str, bytes] = {}
    key_to_hash: dict[bytes, str] = {}
    hash_to_key_conflict_count = 0
    key_to_hash_conflict_count = 0
    total_tokens = 0
    full_blocks = 0
    partial_blocks = 0
    length_digest = hashlib.sha256()
    token_digest = hashlib.sha256()
    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 offset, (row, token_ids) in enumerate(zip(batch, encoded, strict=True)):
            row_index = start + offset
            actual_length = len(token_ids)
            expected_length = int(row["input_length"])
            total_tokens += actual_length
            length_digest.update(f"{row_index}:{actual_length}\n".encode())
            if actual_length != expected_length:
                length_mismatch_count += 1

            source_hashes = row["hash_ids"]
            expected_hashes = math.ceil(actual_length / args.source_block_size)
            if len(source_hashes) != expected_hashes:
                hash_count_mismatch_count += 1
                continue

            request_token_digest = hashlib.sha256()
            parent = b"ROOT"
            for block_index, source_hash in enumerate(source_hashes):
                begin = block_index * args.source_block_size
                chunk = token_ids[begin : begin + args.source_block_size]
                token_payload = b"".join(
                    int(token_id).to_bytes(4, "little", signed=False)
                    for token_id in chunk
                )
                request_token_digest.update(token_payload)
                chunk_digest = hashlib.blake2b(token_payload, digest_size=16).digest()
                key_digest = hashlib.blake2b(
                    parent + b"\0" + chunk_digest, digest_size=16
                ).digest()
                if len(chunk) == args.source_block_size:
                    full_blocks += 1
                else:
                    partial_blocks += 1

                hash_id = str(source_hash)
                previous_key = hash_to_key.setdefault(hash_id, key_digest)
                if previous_key != key_digest:
                    hash_to_key_conflict_count += 1
                previous_hash = key_to_hash.setdefault(key_digest, hash_id)
                if previous_hash != hash_id:
                    key_to_hash_conflict_count += 1
                parent = hash_id.encode()

            token_digest.update(row_index.to_bytes(4, "little"))
            token_digest.update(request_token_digest.digest())

    payload: dict[str, Any] = {
        "schema": "qwen235b-trace-contract-audit-v1",
        "status": "pass_offline_source_contract"
        if not any(
            (
                length_mismatch_count,
                hash_count_mismatch_count,
                hash_to_key_conflict_count,
                key_to_hash_conflict_count,
            )
        )
        else "fail",
        "execution": {
            "host": socket.gethostname(),
            "device": "cpu_only",
            "elapsed_seconds": round(time.time() - started, 3),
            "python_version": platform.python_version(),
            "tokenizer_class": type(tokenizer).__name__,
            "transformers_version": transformers.__version__,
            "model_path": str(args.model.resolve()),
        },
        "trace": {
            "path": str(args.trace.resolve()),
            "sha256": sha256_file(args.trace),
            "source_request_count": len(rows),
            "context_exceeded_count": len(context_exceeded),
            "zero_output_count": len(zero_output),
            "exclusion_overlap_count": sum(row in zero_output for row in context_exceeded),
            "eligible_request_count": len(eligible),
        },
        "tokenization": {
            "total_token_count": total_tokens,
            "input_length_mismatch_count": length_mismatch_count,
            "length_order_sha256": length_digest.hexdigest(),
            "per_request_token_digest_sha256": token_digest.hexdigest(),
        },
        "source_hash_contract": {
            "source_block_size_tokens": args.source_block_size,
            "hash_count_mismatch_count": hash_count_mismatch_count,
            "full_block_count": full_blocks,
            "partial_block_count": partial_blocks,
            "unique_hash_id_count": len(hash_to_key),
            "unique_parent_chunk_key_count": len(key_to_hash),
            "hash_id_to_parent_chunk_conflict_count": hash_to_key_conflict_count,
            "parent_chunk_to_hash_id_conflict_count": key_to_hash_conflict_count,
            "key_definition": (
                "(parent source hash id, BLAKE2b-128 of the tokenizer token-id chunk)"
            ),
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    if payload["status"] != "pass_offline_source_contract":
        raise SystemExit(1)


if __name__ == "__main__":
    main()
