"""Bundle loading, seeded subsets, length-bucketed steps and token-budget micro-batches (pure Python).

One optimiser step takes ``posts_per_step`` posts (classifier-plan §4.3: 32). Steps are drawn from length-sorted
pools so their posts have similar lengths, then shuffled; each step is split into micro-batches whose padded size
(rows x longest row) stays within the token budget (16k). The loss is normalised per step (``step_counts``), so how a
step is split into micro-batches does not change its gradient. No packing.
"""
from __future__ import annotations

import hashlib
import json
import os
import random
from pathlib import Path
from typing import Iterable

SPLITS = ("train", "val", "test")


def load_jsonl(path) -> list[dict]:
    with open(path, encoding="utf-8") as fh:
        return [json.loads(line) for line in fh if line.strip()]


def write_jsonl(path, rows: Iterable[dict], mode: int = 0o600) -> None:
    """Atomic JSONL write (temp file + rename), file mode 0600 by default (rows may hold post text)."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + ".tmp")
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode)
    with os.fdopen(fd, "w", encoding="utf-8") as fh:
        for row in rows:
            fh.write(json.dumps(row, ensure_ascii=False, allow_nan=False, sort_keys=True) + "\n")
    os.chmod(tmp, mode)
    tmp.replace(path)


def load_split(data_dir, split: str) -> list[dict]:
    return load_jsonl(Path(data_dir) / f"{split}.jsonl")


def seeded_order(ids, seed) -> list:
    return sorted(ids, key=lambda i: hashlib.sha256(json.dumps([str(seed), i]).encode()).hexdigest())


def subset(rows: list[dict], n: int, seed) -> list[dict]:
    """The first n rows in seeded hash order of their ids (independent of file order); all rows if n >= len."""
    if not n or n >= len(rows):
        return list(rows)
    keep = set(seeded_order([r["id"] for r in rows], seed)[:n])
    return [r for r in rows if r["id"] in keep]


def steps(lengths: list[int], posts_per_step: int, seed: int, pool: int = 64) -> list[list[int]]:
    """Shuffled steps of similar-length posts: shuffle, sort within pools of posts_per_step*pool, chunk, shuffle."""
    rng = random.Random(seed)
    idx = list(range(len(lengths)))
    rng.shuffle(idx)
    out = []
    for s in range(0, len(idx), posts_per_step * pool):
        chunk = sorted(idx[s:s + posts_per_step * pool], key=lambda i: (lengths[i], i))
        out += [chunk[i:i + posts_per_step] for i in range(0, len(chunk), posts_per_step)]
    rng.shuffle(out)
    return out


def micro_batches(step: list[int], lengths: list[int], token_budget: int) -> list[list[int]]:
    """Split one step so each micro-batch's padded size (rows x longest) is at most token_budget."""
    order = sorted(step, key=lambda i: (-lengths[i], i))
    out: list[list[int]] = []
    current: list[int] = []
    for i in order:
        width = max([lengths[j] for j in current] + [lengths[i]])
        if current and width * (len(current) + 1) > token_budget:
            out.append(current)
            current = []
        current.append(i)
    if current:
        out.append(current)
    return out


def step_counts(labels: list[list[int]]) -> list[int]:
    """Unmasked targets per field over one step's label rows (-1 marks a masked target); at least 1 per field."""
    if not labels:
        return []
    return [max(1, sum(1 for row in labels if row[f] >= 0)) for f in range(len(labels[0]))]


def inference_batches(lengths: list[int], token_budget: int, max_rows: int = 0) -> list[list[int]]:
    """Length-sorted batches for prediction, each within the padded token budget (and max_rows if given)."""
    out: list[list[int]] = []
    current: list[int] = []
    for i in sorted(range(len(lengths)), key=lambda i: (lengths[i], i)):
        full = max_rows and len(current) >= max_rows
        if current and (full or lengths[i] * (len(current) + 1) > token_budget):
            out.append(current)
            current = []
        current.append(i)
    if current:
        out.append(current)
    return out