//! CSM generation CLI.
//!
//! Downloads CSM-1B + Mimi + Llama tokenizer from HuggingFace (cached), then
//! generates speech for the given `--text`.
//!
//! Usage:
//! ```
//! cargo run -p rtx-csm --release --example generate -- \
//!     --text "Hello from Rust." --speaker 0 --out /tmp/hello.wav
//! ```
//!
//! Context conditioning (optional, from a prior utterance you have WAV for).
//! `--context-wav` and `--context-text` are repeatable and zipped in order:
//! ```
//! cargo run -p rtx-csm --release --example generate -- \
//!     --text "Nice to talk to you." \
//!     --context-wav prior_a.wav --context-text "First clip transcript." \
//!     --context-wav prior_b.wav --context-text "Second clip transcript." \
//!     --context-speaker 1 \
//!     --speaker 0 --out /tmp/reply.wav
//! ```
//! Tip: rank manifest clips for context suitability with `scripts/pick_context.sh`.

use anyhow::Result;
use clap::Parser;
use rtx_csm::{GenerateOptions, Generator, PostProcess, Segment, audio_io};

#[derive(Debug, Parser)]
#[command(name = "csm-generate", about = "Generate speech with rtx-csm")]
struct Cli {
    /// Text to synthesize.
    #[arg(long)]
    text: String,

    /// Speaker id (0 or 1).
    #[arg(long, default_value_t = 0)]
    speaker: u32,

    /// Output WAV file (24 kHz mono 16-bit).
    #[arg(long)]
    out: std::path::PathBuf,

    /// Max audio length in milliseconds.
    #[arg(long, default_value_t = 10_000)]
    max_audio_ms: u32,

    /// Sampling temperature.
    #[arg(long, default_value_t = 0.9)]
    temperature: f64,

    /// Top-K sampling.
    #[arg(long, default_value_t = 50)]
    top_k: usize,

    /// Top-p (nucleus) cutoff applied after top-k. 1.0 disables.
    #[arg(long, default_value_t = 0.9)]
    top_p: f64,

    /// RNG seed.
    #[arg(long, default_value_t = 42)]
    seed: u64,

    /// Optional emotion control tag prepended to the text prompt, e.g.
    /// `--emotion-hint "[whisper]"`. Only meaningful when the model has been
    /// fine-tuned with matching tags in its training prompts (see
    /// `docs/personal_voice_training_guide.md`); on the un-adapted base it is
    /// just an extra cosmetic prefix the BPE tokenizer encodes as ordinary
    /// tokens.
    #[arg(long)]
    emotion_hint: Option<String>,

    /// Prior utterance audio (WAV) for context conditioning. Repeatable;
    /// each `--context-wav` must be paired with a matching `--context-text`
    /// in the same order.
    #[arg(long)]
    context_wav: Vec<std::path::PathBuf>,

    /// Transcription of the context WAV. Repeatable; must have the same
    /// count as `--context-wav`.
    #[arg(long)]
    context_text: Vec<String>,

    /// Speaker id assigned to all context utterances. CSM only distinguishes
    /// 0/1; pick something different from `--speaker` if you want the model
    /// to treat the cloning reference as a separate voice.
    #[arg(long, default_value_t = 1)]
    context_speaker: u32,

    /// Force CPU device even if cuda/metal features are enabled.
    #[arg(long)]
    cpu: bool,

    /// Load a quantized GGUF model instead of the default safetensors path.
    /// Path should point at the file produced by `examples/quantize`.
    #[arg(long)]
    quantized_gguf: Option<std::path::PathBuf>,

    /// Apply a LoRA adapter (safetensors) trained via `examples/lora_train`.
    /// Adapter is injected into the FP backbone before generation.
    #[arg(long)]
    lora: Option<std::path::PathBuf>,

    /// LoRA rank override. When omitted, auto-detected from the adapter
    /// file's embedded metadata (Phase 12.5); falls back to 8 for older
    /// adapters without metadata.
    #[arg(long)]
    lora_rank: Option<usize>,

    /// LoRA alpha override. When omitted, auto-detected from the adapter's
    /// embedded metadata; falls back to 16.0 for older adapters.
    #[arg(long)]
    lora_alpha: Option<f32>,

    /// Force extended LoRA coverage (q+k+v+output_proj + MLP). When omitted,
    /// auto-detected from the adapter's embedded metadata. Setting this on a
    /// classic q+v adapter just allocates extra unused B=0 slots.
    #[arg(long, default_value_t = false)]
    extended_lora: bool,

    /// Activation-steering safetensors (one (embed_dim,) tensor per layer
    /// keyed `layer_<i>_steering`). Vectors are added to each backbone
    /// layer's residual stream output during generation. See
    /// `crate::steering` and Sprint 2 of the Phase 8 roadmap. Quantized
    /// GGUF backbone is not supported.
    #[arg(long)]
    steering_vec: Option<std::path::PathBuf>,

    /// Global multiplier on every steering vector. EmoSteer-TTS uses 2.0
    /// for emotion conversion, 2.5 for erasure. 0.0 disables steering
    /// without dropping the file.
    #[arg(long, default_value_t = 1.0)]
    steering_scale: f32,

    /// Comma-separated layer indices to actually steer (e.g. "4,8,12").
    /// Defaults to all layers in the loaded safetensors. EmoSteer-TTS
    /// found that perturbing a small spaced subset of middle-to-deep
    /// layers preserves fluency better than perturbing every layer.
    #[arg(long, value_delimiter = ',')]
    steering_layers: Option<Vec<usize>>,

    /// Steering vectors for the depth decoder (acoustic codebooks).
    /// Decoder is Llama100M for CSM-1B → 4 layers × 1024 embed_dim, so
    /// these are SMALLER than backbone vectors. Architectural hypothesis:
    /// decoder steering shifts prosody/timbre without disturbing word
    /// content the way backbone steering does.
    #[arg(long)]
    decoder_steering_vec: Option<std::path::PathBuf>,

    /// Scale on the decoder steering vector. Default 1.0.
    #[arg(long, default_value_t = 1.0)]
    decoder_steering_scale: f32,

    /// Comma-separated layer indices to actually steer on the decoder
    /// (e.g. "2,3"). The decoder has 4 layers (CSM-1B); useful subsets
    /// to try: `3` (deepest), `2,3` (last two), `1,2` (middle two).
    /// Defaults to all 4 layers.
    #[arg(long, value_delimiter = ',')]
    decoder_steering_layers: Option<Vec<usize>>,

    /// Per-frame CFG scale schedule (Selective CFG, arXiv 2509.19668).
    /// Forms:
    ///   - `const:<scale>`               (constant scale, equivalent to --cfg-scale)
    ///   - `step:<early>:<late>:<T>`     (early until frame T, then late)
    ///   - `linear:<start>:<end>:<R>`    (interpolate over R frames, hold)
    ///
    /// Requires non-empty context AND CFG-enabled load (--enable-cfg).
    /// Recommended: step:3.0:1.5:12 — full CFG for the first ~1s, then
    /// drop to preserve text adherence.
    #[arg(long)]
    cfg_schedule: Option<String>,

    /// Constant CFG scale (legacy single-value form). Ignored if
    /// `--cfg-schedule` is also passed.
    #[arg(long)]
    cfg_scale: Option<f64>,

    /// Load with the CFG-capable dual backbone. Required for any cfg_*
    /// flag to take effect. Adds ~1.5x memory; only worth it when you're
    /// actually using CFG.
    #[arg(long, default_value_t = false)]
    enable_cfg: bool,

    /// Disable audio post-processing (HPF + declick + LUFS normalize).
    #[arg(long)]
    raw: bool,

    /// LUFS target for loudness normalization. Ignored if --raw.
    #[arg(long, default_value_t = -16.0)]
    lufs: f32,

    /// Path to converted AudioSeal generator safetensors (run
    /// `audioseal_convert` first). When set together with
    /// `--watermark-detector`, the watermarker is wired into
    /// `generate_to_wav` so output is automatically watermarked.
    #[arg(long)]
    watermark_generator: Option<std::path::PathBuf>,

    /// Path to converted AudioSeal detector safetensors. Required for the
    /// watermarker even if you only want to embed (the detector is part of
    /// AudioSealWatermarker construction; future builds may make it
    /// optional).
    #[arg(long)]
    watermark_detector: Option<std::path::PathBuf>,

    /// 16-bit watermark message (decimal or 0xHEX).
    #[arg(long, default_value = "0")]
    watermark_message: String,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();

    let cli = Cli::parse();
    let device = if cli.cpu {
        candle_core::Device::Cpu
    } else {
        Generator::default_device()?
    };
    tracing::info!("device: {device:?}");

    let mut generator = if let Some(gguf) = cli.quantized_gguf.as_ref() {
        tracing::info!("loading quantized CSM from {}", gguf.display());
        Generator::load_csm_1b_quantized(gguf, &device, false)?
    } else if cli.enable_cfg {
        tracing::info!("loading CSM-1B with CFG-enabled dual backbone");
        Generator::load_csm_1b_with_cfg(&device, true)?
    } else {
        Generator::load_csm_1b(&device)?
    };

    if let Some(steering_path) = cli.steering_vec.as_ref() {
        let n = generator
            .backbone_num_layers()
            .ok_or_else(|| anyhow::anyhow!("steering not supported on the quantized backbone"))?;
        let mut steering =
            rtx_csm::steering::LayerSteering::load_safetensors(steering_path, n, &device)?;
        steering.set_scale(cli.steering_scale);
        if let Some(layers) = cli.steering_layers.as_ref() {
            steering.restrict_to_layers(layers);
        }
        tracing::info!(
            "loaded backbone steering from {} (scale={}, active_layers={:?})",
            steering_path.display(),
            cli.steering_scale,
            steering.active_layers(),
        );
        generator.set_steering(Some(steering))?;
    }

    if let Some(dec_path) = cli.decoder_steering_vec.as_ref() {
        let n = generator
            .decoder_num_layers()
            .ok_or_else(|| anyhow::anyhow!("decoder steering needs the FP backbone"))?;
        let mut steering =
            rtx_csm::steering::LayerSteering::load_safetensors(dec_path, n, &device)?;
        steering.set_scale(cli.decoder_steering_scale);
        if let Some(layers) = cli.decoder_steering_layers.as_ref() {
            steering.restrict_to_layers(layers);
        }
        tracing::info!(
            "loaded decoder steering from {} (scale={}, active_layers={:?})",
            dec_path.display(),
            cli.decoder_steering_scale,
            steering.active_layers(),
        );
        generator.set_decoder_steering(Some(steering))?;
    }

    if let Some(lora_path) = cli.lora.as_ref() {
        // `--extended-lora` is a force flag (no way to clap-distinguish "off"
        // from "unset" on a bare bool). When the user passes it we override
        // to true; otherwise we let the file metadata or the default decide.
        let extended_override = if cli.extended_lora { Some(true) } else { None };
        rtx_csm::training::apply_lora_adapter(
            &mut generator,
            lora_path,
            cli.lora_rank,
            cli.lora_alpha,
            extended_override,
            &device,
        )?;
    }
    tracing::info!(
        "model loaded (sr={} Hz, frame_rate={} Hz, codebooks={})",
        generator.config.sample_rate,
        generator.config.frame_rate_hz,
        generator.config.audio_num_codebooks,
    );

    if cli.context_wav.len() != cli.context_text.len() {
        anyhow::bail!(
            "--context-wav and --context-text must be repeated the same number of times \
             (got {} wav(s), {} text(s))",
            cli.context_wav.len(),
            cli.context_text.len(),
        );
    }
    let mut context: Vec<Segment> = Vec::with_capacity(cli.context_wav.len());
    for (wav, txt) in cli.context_wav.iter().zip(cli.context_text.iter()) {
        let audio = audio_io::load_mono_24k(wav)?;
        tracing::info!(
            "loaded context #{}: {} samples ({:.2} s) from {}",
            context.len(),
            audio.len(),
            audio.len() as f32 / generator.config.sample_rate as f32,
            wav.display()
        );
        context.push(Segment::new(cli.context_speaker, txt, audio));
    }

    // Optional watermarker wiring (AudioSeal + 24k↔16k resample adapter).
    if let (Some(gen_path), Some(det_path)) = (
        cli.watermark_generator.as_ref(),
        cli.watermark_detector.as_ref(),
    ) {
        let msg_str = cli.watermark_message.trim();
        let message: u16 = if let Some(rest) = msg_str
            .strip_prefix("0x")
            .or_else(|| msg_str.strip_prefix("0X"))
        {
            u16::from_str_radix(rest, 16)?
        } else {
            msg_str.parse::<u16>()?
        };
        let gen_vb = unsafe {
            candle_nn::VarBuilder::from_mmaped_safetensors(
                &[gen_path],
                candle_core::DType::F32,
                &device,
            )
        }?;
        let det_vb = unsafe {
            candle_nn::VarBuilder::from_mmaped_safetensors(
                &[det_path],
                candle_core::DType::F32,
                &device,
            )
        }?;
        let inner = rtx_csm::AudioSealWatermarker::from_var_builders(
            gen_vb,
            det_vb,
            device.clone(),
            message,
        )?;
        // CSM produces 24 kHz; AudioSeal native is 16 kHz.
        let wm = rtx_csm::ResampledWatermarker::new(inner, generator.config.sample_rate, 16_000);
        generator.set_watermarker(Box::new(wm));
        tracing::info!(
            "watermarker installed (message=0x{:04X}, model 16 kHz, output {} Hz)",
            message,
            generator.config.sample_rate
        );
    }

    let cfg_schedule = match (cli.cfg_schedule.as_ref(), cli.cfg_scale) {
        (Some(s), _) => Some(rtx_csm::cfg_schedule::CfgSchedule::parse(s)?),
        (None, Some(_)) => None, // legacy cfg_scale path handled by GenerateOptions::cfg_scale
        (None, None) => None,
    };
    let opts = GenerateOptions {
        max_audio_ms: cli.max_audio_ms,
        temperature: cli.temperature,
        top_k: cli.top_k,
        top_p: cli.top_p,
        seed: cli.seed,
        emotion_hint: cli.emotion_hint.clone(),
        cfg_scale: cli.cfg_scale,
        cfg_schedule,
        ..GenerateOptions::default()
    };
    let post = if cli.raw {
        PostProcess::disabled()
    } else {
        PostProcess {
            lufs_target: Some(cli.lufs),
            ..PostProcess::default()
        }
    };
    generator.generate_to_wav(&cli.text, cli.speaker, &context, opts, &post, &cli.out)?;
    println!("wrote {}", cli.out.display());
    Ok(())
}