//! Generate the M1 verifiable-arithmetic post-training dataset. Pure host tool (no
//! CUDA): writes
//!   <out>/arith_sft.tsv           user<TAB>assistant rows for `train --sft-tsv`
//!   <out>/arith_eval_prompts.txt  greedy_sample `--prompts-file` format (held out)
//!   <out>/arith_eval_gold.txt     parallel gold integers for the checker
//!
//! Eval problems are deduped against train (no leakage). The SFT rows carry just the
//! user/assistant content; `data::load_sft_tsv_cached` adds the `User:/Assistant:`
//! frame + `<|endoftext|>` and masks the prompt, so the eval prompt lines here
//! reconstruct exactly that frame (`User: <q>\nAssistant:`, literal `\n` decoded by
//! greedy_sample).
//!
//! Example:
//!   cargo run -p xtrain-train --release --bin gen_arith_task -- \
//!       --n 20000 --eval 500 --seed 1 --out-dir /dashscope-tmp/wjh/xtrain_post/arith

use std::collections::HashSet;
use std::fs::{self, File};
use std::io::{BufWriter, Write};
use std::path::PathBuf;

use xtrain_train::task::{GenConfig, Op, gen_problem};

fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
    args.iter()
        .position(|a| a == name)
        .and_then(|i| args.get(i + 1))
        .and_then(|v| v.parse().ok())
        .unwrap_or(default)
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let n_train: usize = flag(&args, "--n", 20000);
    let n_eval: usize = flag(&args, "--eval", 500);
    let seed: u64 = flag(&args, "--seed", 1);
    let max_add: i64 = flag(&args, "--max-add", 99);
    let max_mul: i64 = flag(&args, "--max-mul", 12);
    let out_dir: PathBuf = args
        .iter()
        .position(|a| a == "--out-dir")
        .and_then(|i| args.get(i + 1))
        .map(PathBuf::from)
        .expect("--out-dir <dir> is required");

    fs::create_dir_all(&out_dir).expect("create out dir");
    let cfg = GenConfig {
        max_add,
        max_mul,
        ops: vec![Op::Add, Op::Sub, Op::Mul],
    };
    let mut rng = seed.max(1);

    // Train: dedup so the same problem is not repeated and so eval can be held out.
    let mut train_keys = HashSet::new();
    let mut tsv = BufWriter::new(File::create(out_dir.join("arith_sft.tsv")).expect("create tsv"));
    while train_keys.len() < n_train {
        let p = gen_problem(&mut rng, &cfg);
        if !train_keys.insert(p.key()) {
            continue;
        }
        writeln!(tsv, "{}\t{}", p.question(), p.sft_answer()).expect("write tsv");
    }
    tsv.flush().expect("flush tsv");

    // Eval: disjoint from train (skip any key seen in train) and from itself.
    let mut prompts =
        BufWriter::new(File::create(out_dir.join("arith_eval_prompts.txt")).expect("create eval"));
    let mut golds =
        BufWriter::new(File::create(out_dir.join("arith_eval_gold.txt")).expect("create gold"));
    writeln!(prompts, "# verifiable arithmetic eval prompts (held out from arith_sft.tsv)")
        .expect("write header");
    let mut eval_keys = HashSet::new();
    while eval_keys.len() < n_eval {
        let p = gen_problem(&mut rng, &cfg);
        if train_keys.contains(&p.key()) || !eval_keys.insert(p.key()) {
            continue;
        }
        writeln!(prompts, "User: {}\\nAssistant:", p.question()).expect("write prompt");
        writeln!(golds, "{}", p.answer()).expect("write gold");
    }
    prompts.flush().expect("flush prompts");
    golds.flush().expect("flush golds");

    println!(
        "wrote {} train rows + {} eval prompts to {} (ops=+,-,* max_add={} max_mul={} seed={})",
        train_keys.len(),
        eval_keys.len(),
        out_dir.display(),
        max_add,
        max_mul,
        seed
    );
}