//! End-to-end production pipeline showcase.
//!
//! Single command demonstrates the full rtx-csm stack:
//!   1. CSM-1B generates 24 kHz speech from text (with optional LoRA voice).
//!   2. Post-process: HPF + declick + EBU R128 −16 LUFS.
//!   3. AudioSeal watermark embedded transparently via the resampling
//!      adapter (24 kHz ↔ 16 kHz round-trip).
//!   4. AudioSeal detector verifies the watermark on the written WAV.
//!   5. WavLM-SV computes a 512-d speaker embedding for the output and,
//!      if a reference WAV is provided, scores cosine similarity against it.
//!
//! Usage:
//! ```
//! cargo run -p rtx-csm --release --features metal --example pipeline -- \
//!     --text "Hello from the full Rust pipeline." \
//!     --speaker 0 \
//!     --audioseal-generator /tmp/audioseal_generator.safetensors \
//!     --audioseal-detector  /tmp/audioseal_detector.safetensors \
//!     --audioseal-message 0xCAFE \
//!     --wavlm-sv /tmp/wavlm_sv.safetensors \
//!     --reference /tmp/csm_24k.wav \
//!     --out /tmp/pipeline_out.wav
//! ```
//!
//! The reference WAV is expected to be a known-speaker sample (any rate;
//! resampled to 16 kHz internally). Cosine sim ≈ 1.0 confirms the output
//! sounds like the same speaker as the reference; lower values flag drift.

use anyhow::Result;
use candle_core::{DType, Device};
use clap::Parser;
use rtx_csm::{
    GenerateOptions, Generator, PostProcess, Segment, audio_io,
    audioseal::AudioSealWatermarker,
    speaker_sim::{SpeakerSimilarity, WavLmSimilarity},
    watermark::ResampledWatermarker,
};
use std::path::PathBuf;

const AUDIOSEAL_RATE: u32 = 16_000;

#[derive(Debug, Parser)]
#[command(name = "pipeline", about = "rtx-csm end-to-end showcase")]
struct Cli {
    /// Text to synthesize.
    #[arg(long)]
    text: String,
    /// Speaker id (0 or 1).
    #[arg(long, default_value_t = 0)]
    speaker: u32,
    /// Output WAV path.
    #[arg(long)]
    out: PathBuf,
    /// Max audio length in milliseconds.
    #[arg(long, default_value_t = 8_000)]
    max_audio_ms: u32,
    /// LUFS target for loudness normalization.
    #[arg(long, default_value_t = -16.0)]
    lufs: f32,
    /// Optional LoRA adapter (run examples/lora_train to produce one).
    #[arg(long)]
    lora: Option<PathBuf>,
    #[arg(long, default_value_t = 8)]
    lora_rank: usize,
    #[arg(long, default_value_t = 16.0)]
    lora_alpha: f32,

    /// AudioSeal generator safetensors (produced by audioseal_convert).
    /// If both --audioseal-generator and --audioseal-detector are set,
    /// output is watermarked + detection round-trip is run.
    #[arg(long)]
    audioseal_generator: Option<PathBuf>,
    #[arg(long)]
    audioseal_detector: Option<PathBuf>,
    /// 16-bit watermark message (decimal or 0xHEX).
    #[arg(long, default_value = "0xCAFE")]
    audioseal_message: String,

    /// WavLM-SV safetensors (from wavlm_sv_convert). When set, the
    /// pipeline computes a 512-d speaker embedding for the output WAV.
    #[arg(long)]
    wavlm_sv: Option<PathBuf>,
    /// Optional reference WAV to score the output's speaker against.
    #[arg(long)]
    reference: Option<PathBuf>,

    /// Force CPU device.
    #[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 {
        Device::Cpu
    } else {
        Generator::default_device()?
    };
    println!("== rtx-csm pipeline ==");
    println!("device: {device:?}");

    // -- 1. Build the Generator --------------------------------------------
    let t0 = std::time::Instant::now();
    let mut generator = Generator::load_csm_1b(&device)?;
    println!(
        "[1/5] CSM-1B loaded in {:.2}s (sr={} Hz)",
        t0.elapsed().as_secs_f32(),
        generator.config.sample_rate
    );

    // Optional LoRA voice clone.
    if let Some(lora_path) = cli.lora.as_ref() {
        let lora_cfg = rtx_csm::lora::LoraConfig {
            rank: cli.lora_rank,
            alpha: cli.lora_alpha,
            ..rtx_csm::lora::LoraConfig::default()
        };
        let vm = candle_nn::VarMap::new();
        generator.model.inner.add_lora_to_backbone(&lora_cfg, &vm)?;
        rtx_csm::training::load_lora_adapter(&vm, lora_path, &device)?;
        generator.model.inner.refresh_lora(&vm)?;
        println!(
            "      LoRA adapter loaded: {} (rank={} alpha={})",
            lora_path.display(),
            cli.lora_rank,
            cli.lora_alpha
        );
    }

    // -- 2. Optional inline watermarker -----------------------------------
    let watermark_message = parse_message(&cli.audioseal_message)?;
    let mut audioseal_for_detect: Option<AudioSealWatermarker> = None;
    if let (Some(g), Some(d)) = (
        cli.audioseal_generator.as_ref(),
        cli.audioseal_detector.as_ref(),
    ) {
        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(), watermark_message)?;
        let wm = ResampledWatermarker::new(inner, generator.config.sample_rate, AUDIOSEAL_RATE);
        generator.set_watermarker(Box::new(wm));
        // Build a second instance for detect-only (the one above moves into the
        // generator). This is cheap: weights are mmap'd, only metadata is duplicated.
        let g_vb2 =
            unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[g], DType::F32, &device) }?;
        let d_vb2 =
            unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[d], DType::F32, &device) }?;
        audioseal_for_detect = Some(AudioSealWatermarker::from_var_builders(
            g_vb2,
            d_vb2,
            device.clone(),
            watermark_message,
        )?);
        println!(
            "[2/5] AudioSeal watermarker installed (message=0x{:04X})",
            watermark_message
        );
    } else {
        println!(
            "[2/5] AudioSeal: skipped (pass both --audioseal-generator and --audioseal-detector to enable)"
        );
    }

    // -- 3. Generate --------------------------------------------------------
    let opts = GenerateOptions {
        max_audio_ms: cli.max_audio_ms,
        ..GenerateOptions::default()
    };
    let post = PostProcess {
        lufs_target: Some(cli.lufs),
        ..PostProcess::default()
    };
    let context: Vec<Segment> = Vec::new();
    let t_gen = std::time::Instant::now();
    generator.generate_to_wav(&cli.text, cli.speaker, &context, opts, &post, &cli.out)?;
    let gen_secs = t_gen.elapsed().as_secs_f32();
    println!(
        "[3/5] generated + post-processed{} in {:.2}s -> {}",
        if audioseal_for_detect.is_some() {
            " + watermarked"
        } else {
            ""
        },
        gen_secs,
        cli.out.display()
    );

    // -- 4. Watermark verification round-trip ------------------------------
    if let Some(detector) = audioseal_for_detect.as_ref() {
        let raw = audio_io::load_mono_at_rate(&cli.out, AUDIOSEAL_RATE)?;
        let result = detector.detect(&raw)?;
        let decoded = result.message.unwrap_or(0);
        let bits_match = 16 - (decoded ^ watermark_message).count_ones() as usize;
        println!(
            "[4/5] AudioSeal detect: mean_presence={:.4}, decoded=0x{:04X}, bits={}/16",
            result.mean_presence, decoded, bits_match
        );
    } else {
        println!("[4/5] AudioSeal detect: skipped");
    }

    // -- 5. WavLM-SV speaker embedding + optional reference scoring -------
    if let Some(wavlm_path) = cli.wavlm_sv.as_ref() {
        let scorer = WavLmSimilarity::load(wavlm_path, &device)?;
        let out_samples = audio_io::load_mono_at_rate(&cli.out, AUDIOSEAL_RATE)?;
        let out_emb = scorer.embed(&out_samples)?;
        println!("[5/5] WavLM-SV embedded output (len={})", out_emb.len());
        if let Some(ref_path) = cli.reference.as_ref() {
            let ref_samples = audio_io::load_mono_at_rate(ref_path, AUDIOSEAL_RATE)?;
            let sim = scorer.score(&out_samples, &ref_samples)?;
            println!(
                "      cosine vs reference {}: {:.4}  ({})",
                ref_path.display(),
                sim,
                if sim > 0.5 {
                    "likely same speaker"
                } else {
                    "likely different speakers"
                }
            );
        }
    } else {
        println!("[5/5] WavLM-SV: skipped (pass --wavlm-sv to enable)");
    }

    println!("== pipeline complete: {} ==", cli.out.display());
    Ok(())
}