#!/usr/bin/env python3
"""Build a deterministic SFT mixture containing only complete trajectories.

Inputs are already-decontaminated, public trajectory datasets.  Completion is
defined structurally and without reference to evaluation tasks: the final
message must be a non-tool assistant response.  This excludes source rollouts
that ended at their turn limit or immediately after a shell observation.
"""

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 is_complete(messages: list[dict]) -> bool:
    if not messages:
        return False
    final = messages[-1]
    return (
        final.get("role") == "assistant"
        and not final.get("tool_calls")
        and bool((final.get("content") or "").strip())
    )


def component(path: Path, label: str, num_proc: int) -> tuple[Dataset, int]:
    dataset = load_dataset(str(path), split="train").select_columns(COMMON_COLUMNS)
    input_rows = len(dataset)
    dataset = dataset.filter(
        is_complete,
        input_columns=["messages"],
        num_proc=num_proc,
        desc=f"Keeping complete {label} trajectories",
    )
    dataset = dataset.add_column("mix_component", [label] * len(dataset))
    return dataset, input_rows


def write_shards(dataset: Dataset, output_dir: Path, shards: int) -> None:
    output_dir.mkdir(parents=True, exist_ok=False)
    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("--output", type=Path, required=True)
    parser.add_argument("--terminal-repeats", type=int, default=2)
    parser.add_argument("--swe-repeats", type=int, default=2)
    parser.add_argument("--seed", type=int, default=20260814)
    parser.add_argument("--shards", type=int, default=64)
    parser.add_argument("--num-proc", type=int, default=32)
    args = parser.parse_args()

    terminal, terminal_input_rows = component(
        args.terminal, "nemotron-terminal-skill-complete", args.num_proc
    )
    swe, swe_input_rows = component(args.swe, "swe-smith-complete", args.num_proc)
    parts = [terminal] * args.terminal_repeats + [swe] * args.swe_repeats
    mixed = concatenate_datasets(parts).shuffle(seed=args.seed)
    write_shards(mixed, args.output, min(args.shards, len(mixed)))

    report = {
        "selection": "final message is a non-tool assistant response with non-empty content",
        "selection_uses_evaluation_content": False,
        "terminal_path": str(args.terminal),
        "terminal_input_rows": terminal_input_rows,
        "terminal_complete_rows": len(terminal),
        "terminal_repeats": args.terminal_repeats,
        "swe_path": str(args.swe),
        "swe_input_rows": swe_input_rows,
        "swe_complete_rows": len(swe),
        "swe_repeats": args.swe_repeats,
        "output_rows": len(mixed),
        "shuffle_seed": args.seed,
        "mix": {
            "nemotron-terminal-skill-complete": len(terminal) * args.terminal_repeats,
            "swe-smith-complete": len(swe) * args.swe_repeats,
        },
    }
    (args.output / "provenance.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()