"""Token shards (uint16 .bin files), train/val split and the corpus manifest."""

from __future__ import annotations

import hashlib
import json
from collections.abc import Iterable
from datetime import datetime, timezone
from pathlib import Path
from typing import Literal

import numpy as np


def split_of(record_id: str) -> Literal["train", "val"]:
    return "val" if int(hashlib.sha256(record_id.encode()).hexdigest()[:8], 16) % 200 == 0 else "train"


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


class ShardWriter:
    """Appends `[bos] + ids + [eos]` documents; starts a new shard after `shard_tokens`."""

    def __init__(self, out_dir: Path, prefix: str, bos_id: int, eos_id: int, shard_tokens: int = 100_000_000):
        self.out_dir = Path(out_dir)
        self.out_dir.mkdir(parents=True, exist_ok=True)
        self.prefix, self.bos_id, self.eos_id, self.shard_tokens = prefix, bos_id, eos_id, shard_tokens
        self.tokens = 0
        self.bytes = 0
        self._buf: list[np.ndarray] = []
        self._buf_tokens = 0
        self._shards: list[dict] = []

    def add(self, ids: list[int], n_bytes: int = 0) -> None:
        arr = np.asarray([self.bos_id, *ids, self.eos_id], dtype="<u2")
        self._buf.append(arr)
        self._buf_tokens += len(arr)
        self.tokens += len(arr)
        self.bytes += n_bytes
        if self._buf_tokens >= self.shard_tokens:
            self._flush()

    def _flush(self) -> None:
        if not self._buf:
            return
        path = self.out_dir / f"{self.prefix}_{len(self._shards):03d}.bin"
        np.concatenate(self._buf).tofile(path)
        self._shards.append({"path": str(path), "tokens": self._buf_tokens, "sha256": _sha256(path)})
        self._buf, self._buf_tokens = [], 0

    def close(self) -> list[dict]:
        self._flush()
        return self._shards


def write_shards(token_docs: Iterable[list[int]], out_dir: Path, prefix: str, bos_id: int, eos_id: int,
                 shard_tokens: int = 100_000_000) -> list[dict]:
    w = ShardWriter(out_dir, prefix, bos_id, eos_id, shard_tokens)
    for ids in token_docs:
        w.add(ids)
    return w.close()


def write_manifest(out_dir: Path, stats: dict) -> None:
    manifest = {"created": datetime.now(timezone.utc).isoformat(timespec="seconds"), **stats}
    (Path(out_dir) / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True))