from __future__ import annotations

import argparse
import json
from pathlib import Path

import pandas as pd


ROOT = Path("/home/ubuntu/transcripts")
DEFAULT_INPUT = ROOT / "data" / "transcript_variant_prompt_samples.json"
DEFAULT_OUTPUT = ROOT / "data" / "transcript_variant_test_input.parquet"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Build local test parquet for transcript variant workers.")
    parser.add_argument("--input-json", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--rows", type=int, default=240)
    return parser.parse_args()


def main():
    args = parse_args()
    samples = json.loads(args.input_json.read_text())
    rows = []
    for idx in range(args.rows):
        sample = dict(samples[idx % len(samples)])
        sample["row_id"] = f"variant_row_{idx:06d}"
        sample["video_id"] = f"variant_video_{idx // 12:05d}"
        sample["segment_id"] = f"{sample['video_id']}_seg_{idx:06d}"
        sample["text"] = sample.pop("text")
        rows.append(sample)

    df = pd.DataFrame.from_records(rows)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(args.output, index=False)
    print(f"Wrote {len(df)} rows -> {args.output}")


if __name__ == "__main__":
    main()