#!/usr/bin/env python3
"""Build a deterministic, balanced terminal/SWE SFT mixture."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from datasets import Dataset, concatenate_datasets, load_dataset

COMMON_COLUMNS = ["messages", "tools", "source", "num_tool_calls"]


def normalized_component(path: Path, label: str) -> Dataset:
    dataset = load_dataset(str(path), split="train").select_columns(COMMON_COLUMNS)
    return dataset.add_column("mix_component", [label] * len(dataset))


def write_shards(dataset: Dataset, output_dir: Path, shards: int) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    for index in range(shards):
        dataset.shard(shards, index, contiguous=True).to_parquet(
            output_dir / f"train-{index:05d}-of-{shards:05d}.parquet"
        )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--terminal", type=Path, required=True)
    parser.add_argument("--swe", type=Path, required=True)
    parser.add_argument("--adapters", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--swe-repeats", type=int, default=2)
    parser.add_argument("--adapter-repeats", type=int, default=1)
    parser.add_argument("--seed", type=int, default=20260814)
    parser.add_argument("--shards", type=int, default=64)
    args = parser.parse_args()

    terminal = normalized_component(args.terminal, "nemotron-terminal-skill")
    swe = normalized_component(args.swe, "swe-smith")
    parts = [terminal, *([swe] * args.swe_repeats)]
    adapters = None
    if args.adapters is not None:
        adapters = normalized_component(args.adapters, "nemotron-terminal-adapters")
        parts.extend([adapters] * args.adapter_repeats)
    mixed = concatenate_datasets(parts).shuffle(seed=args.seed)
    write_shards(mixed, args.output, min(args.shards, len(mixed)))

    report = {
        "terminal_path": str(args.terminal),
        "terminal_rows": len(terminal),
        "swe_path": str(args.swe),
        "swe_rows": len(swe),
        "swe_repeats": args.swe_repeats,
        "adapter_path": str(args.adapters) if args.adapters is not None else None,
        "adapter_rows": len(adapters) if adapters is not None else 0,
        "adapter_repeats": args.adapter_repeats if adapters is not None else 0,
        "output_rows": len(mixed),
        "shuffle_seed": args.seed,
        "mix": {
            "nemotron-terminal-skill": len(terminal),
            "swe-smith": len(swe) * args.swe_repeats,
            **(
                {"nemotron-terminal-adapters": len(adapters) * args.adapter_repeats}
                if adapters is not None
                else {}
            ),
        },
    }
    (args.output / "provenance.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()