#!/usr/bin/env python3
"""Wrapper for self-contained LTX-2 preprocessing inside a derivation."""

import argparse
import os
import subprocess
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--training-data-dir", required=True, type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--resolution-buckets", default="960x544x49")
    parser.add_argument("--ltx2-src", required=True, type=Path)
    parser.add_argument("--ltx2-model", required=True, type=Path)
    parser.add_argument("--gemma-model", required=True, type=Path)
    args = parser.parse_args()

    env = os.environ.copy()
    subprocess.run([
        "python", "packages/ltx-trainer/scripts/process_dataset.py",
        str(args.training_data_dir / "ltx2_dataset.json"),
        "--resolution-buckets", args.resolution_buckets,
        "--model-path", str(args.ltx2_model),
        "--text-encoder-path", str(args.gemma_model),
        "--output-dir", str(args.output_dir),
    ], check=True, env=env, cwd=str(args.ltx2_src))


if __name__ == "__main__":
    main()