"""Describe Parquet files of prepared token rows as a prepared-tokens-v1 dataset.

Each file holds one row per conversation, with an `input_ids` and a `loss_mask` list column of
equal length. `loss_mask[i]` is 1 where token `i` is trained on. Other columns are ignored.

The script reads every row, checks it can be trained as given, and writes `prepared.json`: the
tokenizer's repository and commit, the digests of its tokenizer and chat template, and each file's
digest and counts. With `--zip`, it writes the descriptor and the files into one archive to
register:

    python describe_prepared.py --model google/gemma-4-E4B-it --zip train.zip train.parquet
    akka-optimize registry datasets register gemma/train -f train.zip

It needs `transformers` and `pyarrow`, and reads the tokenizer from the Hugging Face Hub.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
import zipfile
from pathlib import Path
from typing import Any

DESCRIPTOR = "prepared.json"
COLUMNS = ("input_ids", "loss_mask")


class RowError(ValueError):
    """A row that cannot be trained as given, named by its file and position."""


def tokenizer_sha256(tokenizer: Any) -> str:
    """The digest of the tokenizer's serialized form, which the trainer compares."""
    backend = getattr(tokenizer, "backend_tokenizer", None)
    if backend is None:
        raise ValueError("the tokenizer has no serialized form to name the rows by")
    return hashlib.sha256(backend.to_str().encode()).hexdigest()


def chat_template_sha256(tokenizer: Any) -> str:
    """The digest of the chat template the rows were rendered with, which the trainer compares."""
    template = tokenizer.chat_template
    if not isinstance(template, str):
        raise ValueError("the tokenizer has no single chat template to name the rows by")
    return hashlib.sha256(template.encode()).hexdigest()


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


def refuse(where: str, before: int, failing: Any, reason: str) -> None:
    """Raises for the first row `failing` marks, numbered from the start of the file."""
    if failing.any():
        raise RowError(f"{where} {before + int(failing.argmax())}: {reason}")


def count(path: Path, vocabulary: int) -> dict[str, int]:
    """The rows, tokens, supervised tokens and longest row of one file, read a row group at a
    time. Refuses a row the trainer would refuse."""
    import numpy as np
    import pyarrow as pa
    import pyarrow.parquet as pq

    parquet = pq.ParquetFile(path)
    for name in COLUMNS:
        # a float would be cast to an integer on the way to a device, training another token
        kind = parquet.schema_arrow.field(name).type
        if not (pa.types.is_list(kind) or pa.types.is_large_list(kind)) or not (
            pa.types.is_integer(kind.value_type)
        ):
            raise ValueError(f"{path.name} holds {name} as {kind}, not a list of integers")
    rows = tokens = supervised = longest = 0
    for group in range(parquet.num_row_groups):
        table = parquet.read_row_group(group, columns=list(COLUMNS))
        ids = table.column("input_ids").combine_chunks()
        mask = table.column("loss_mask").combine_chunks()
        for name, column in (("input_ids", ids), ("loss_mask", mask)):
            # a null row or element is neither a token nor a mark, and the trainer refuses it
            refuse(
                f"{path.name} row",
                rows,
                column.is_null().to_numpy(zero_copy_only=False),
                f"{name} is null",
            )
            if column.values.null_count:
                raise RowError(f"{path.name} holds a null in {name}")
        ends = ids.offsets.to_numpy()
        lengths = np.diff(ends)
        mask_lengths = np.diff(mask.offsets.to_numpy())
        values = ids.values.to_numpy(zero_copy_only=False)
        marks = mask.values.to_numpy(zero_copy_only=False)
        starts = mask.offsets.to_numpy()[:-1]

        where = f"{path.name} row"
        refuse(where, rows, lengths != mask_lengths, "input_ids and loss_mask differ in length")
        refuse(where, rows, lengths < 2, "a row holds at least two tokens")
        # every row has tokens from here on, so a sum from each row's start covers that row alone
        not_a_bit = ((marks != 0) & (marks != 1)).astype(np.int64)
        refuse(
            where,
            rows,
            np.add.reduceat(not_a_bit, starts) > 0,
            "loss_mask holds a value other than 0 and 1",
        )
        refuse(
            where,
            rows,
            marks[starts] != 0,
            "loss_mask marks the first token, which nothing predicts",
        )
        per_row = np.add.reduceat(marks.astype(np.int64), starts)
        refuse(where, rows, per_row == 0, "loss_mask marks no token")
        if values.size and (values.min() < 0 or values.max() >= vocabulary):
            raise ValueError(f"{path.name} holds a token ID outside the tokenizer's {vocabulary}")
        rows += len(lengths)
        tokens += int(lengths.sum())
        supervised += int(per_row.sum())
        longest = max(longest, int(lengths.max()) if len(lengths) else 0)
    if rows == 0:
        raise ValueError(f"{path.name} holds no rows")
    return {"rows": rows, "tokens": tokens, "supervisedTokens": supervised, "longestRow": longest}


def describe(tokenizer: Any, repository: str, revision: str, paths: list[Path]) -> dict[str, Any]:
    """The descriptor for these files, in training order, made with this tokenizer."""
    if len({path.name for path in paths}) != len(paths):
        raise ValueError(
            "each rows file needs a name of its own, since the version holds them side by side"
        )
    return {
        "tokenizer": {"repository": repository, "revision": revision},
        "tokenizerSha256": tokenizer_sha256(tokenizer),
        "chatTemplateSha256": chat_template_sha256(tokenizer),
        "files": [
            {"path": path.name, "sha256": file_sha256(path), **count(path, len(tokenizer))}
            for path in paths
        ],
    }


def write_zip(target: Path, descriptor: dict[str, Any], paths: list[Path]) -> None:
    # stored rather than deflated: Parquet is compressed already
    with zipfile.ZipFile(target, "w", compression=zipfile.ZIP_STORED) as archive:
        archive.writestr(DESCRIPTOR, json.dumps(descriptor, indent=2) + "\n")
        for path in paths:
            archive.write(path, path.name)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model", required=True, help="the base model repository, as on the Hub")
    parser.add_argument(
        "--revision", default="main", help="the branch, tag or commit; recorded as its commit"
    )
    parser.add_argument(
        "--zip", type=Path, help="write the descriptor and the files to this archive"
    )
    parser.add_argument("files", nargs="+", type=Path, help="the rows files, in training order")
    args = parser.parse_args(argv)

    from huggingface_hub import HfApi
    from transformers import AutoTokenizer

    commit = HfApi().model_info(args.model, revision=args.revision).sha
    tokenizer = AutoTokenizer.from_pretrained(args.model, revision=commit)
    if tokenizer.chat_template is None:
        # a model that reads images can keep its template with its processor
        from transformers import AutoProcessor

        tokenizer.chat_template = AutoProcessor.from_pretrained(
            args.model, revision=commit
        ).chat_template
    try:
        descriptor = describe(tokenizer, args.model, commit, args.files)
    except ValueError as refused:
        print(f"describe_prepared: {refused}", file=sys.stderr)
        return 1
    if args.zip:
        write_zip(args.zip, descriptor, args.files)
    else:
        print(json.dumps(descriptor, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())