//! CSM-1B safetensors → quantized GGUF converter.
//!
//! Builds the artifact a (future) forked `csm_quantized.rs` would consume.
//! For now: prints the policy decision report, writes the quantized GGUF.
//!
//! Usage:
//! ```
//! cargo run -p rtx-csm --release --example quantize -- \
//!   --policy q8 --out /tmp/csm-q8.gguf
//! cargo run -p rtx-csm --release --example quantize -- \
//!   --policy q4km --out /tmp/csm-q4km.gguf
//! ```

use anyhow::Result;
use clap::{Parser, ValueEnum};
use rtx_csm::quantize::{QuantPolicy, convert_to_quantized};
use rtx_csm::{hub, model};

#[derive(Debug, Clone, Copy, ValueEnum)]
enum PolicyChoice {
    /// Q8_0 on backbone projections, F16/F32 on heads/embeds/decoder
    Q8,
    /// Q4_K on backbone projections, Q8_0 on decoder, F16/F32 on heads/embeds
    Q4km,
}

#[derive(Debug, Parser)]
#[command(name = "csm-quantize")]
struct Cli {
    /// Path to input safetensors (defaults to HF-cached sesame/csm-1b weights).
    #[arg(long)]
    input: Option<std::path::PathBuf>,

    /// Output GGUF file.
    #[arg(long)]
    out: std::path::PathBuf,

    /// Policy preset.
    #[arg(long, value_enum, default_value = "q8")]
    policy: PolicyChoice,

    /// Print policy decisions only — don't write anything.
    #[arg(long)]
    dry_run: bool,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();

    let input_path = match cli.input {
        Some(p) => p,
        None => hub::resolve_csm_weights()?,
    };
    eprintln!("input:  {}", input_path.display());
    eprintln!("output: {}", cli.out.display());

    let policy = match cli.policy {
        PolicyChoice::Q8 => QuantPolicy::q8_safe(),
        PolicyChoice::Q4km => QuantPolicy::q4km_aggressive(),
    };

    if cli.dry_run {
        let descs = model::dump_safetensors_keys(&input_path)?;
        let report = rtx_csm::quantize::report(&descs, &policy);
        println!("{report}");
        return Ok(());
    }

    eprintln!("converting (this can take 30-90 seconds for CSM-1B)...");
    let report = convert_to_quantized(&input_path, &cli.out, &policy)?;
    println!("\n{report}");
    Ok(())
}