//! Long-form generation with rolling context, post-processing, and
//! optional inline watermark — exercises `Generator::generate_long_to_wav`.
//!
//! Usage:
//! ```
//! cargo run -p rtx-csm --release --features metal --example generate_long -- \
//!     --text "Long passage..." \
//!     --out /tmp/longform.wav \
//!     --watermark-generator /tmp/audioseal_generator.safetensors \
//!     --watermark-detector  /tmp/audioseal_detector.safetensors
//! ```

use anyhow::Result;
use candle_core::DType;
use clap::Parser;
use rtx_csm::{
    GenerateOptions, Generator, PostProcess, audioseal::AudioSealWatermarker,
    longform::LongFormConfig, watermark::ResampledWatermarker,
};
use std::path::PathBuf;

const AUDIOSEAL_RATE: u32 = 16_000;

#[derive(Debug, Parser)]
#[command(name = "generate_long")]
struct Cli {
    #[arg(long)]
    text: String,
    #[arg(long, default_value_t = 0)]
    speaker: u32,
    #[arg(long)]
    out: PathBuf,
    #[arg(long, default_value_t = 12_000)]
    max_audio_ms: u32,
    /// Max characters per chunk (sentence packer target).
    #[arg(long, default_value_t = 220)]
    max_chunk_chars: usize,
    #[arg(long, default_value_t = -16.0)]
    lufs: f32,
    /// Optional AudioSeal watermarker (both flags required to enable).
    #[arg(long)]
    watermark_generator: Option<PathBuf>,
    #[arg(long)]
    watermark_detector: Option<PathBuf>,
    #[arg(long, default_value = "0xCAFE")]
    watermark_message: String,
    #[arg(long)]
    cpu: bool,
}

fn parse_message(s: &str) -> Result<u16> {
    let s = s.trim();
    if let Some(rest) = s.strip_prefix("0x").or_else(|| s.strip_prefix("0X")) {
        Ok(u16::from_str_radix(rest, 16)?)
    } else {
        Ok(s.parse::<u16>()?)
    }
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();
    let device = if cli.cpu {
        candle_core::Device::Cpu
    } else {
        Generator::default_device()?
    };
    let mut generator = Generator::load_csm_1b(&device)?;
    println!("loaded CSM-1B (sr={} Hz)", generator.config.sample_rate);

    if let (Some(g), Some(d)) = (
        cli.watermark_generator.as_ref(),
        cli.watermark_detector.as_ref(),
    ) {
        let message = parse_message(&cli.watermark_message)?;
        let g_vb =
            unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[g], DType::F32, &device) }?;
        let d_vb =
            unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[d], DType::F32, &device) }?;
        let inner = AudioSealWatermarker::from_var_builders(g_vb, d_vb, device.clone(), message)?;
        let wm = ResampledWatermarker::new(inner, generator.config.sample_rate, AUDIOSEAL_RATE);
        generator.set_watermarker(Box::new(wm));
        println!("watermarker installed (message=0x{:04X})", message);
    }

    let opts = GenerateOptions {
        max_audio_ms: cli.max_audio_ms,
        ..GenerateOptions::default()
    };
    let post = PostProcess {
        lufs_target: Some(cli.lufs),
        ..PostProcess::default()
    };
    let cfg = LongFormConfig {
        max_chunk_chars: cli.max_chunk_chars,
        ..LongFormConfig::default()
    };

    let t = std::time::Instant::now();
    generator.generate_long_to_wav(&cli.text, cli.speaker, None, opts, cfg, &post, &cli.out)?;
    println!(
        "wrote {} in {:.2}s",
        cli.out.display(),
        t.elapsed().as_secs_f32()
    );
    Ok(())
}