//! Watermarking hook for the TTS pipeline.
//!
//! The trait is sample-rate-agnostic: callers feed PCM in whatever native
//! rate they have and receive PCM at the same rate. Wrap a model-bound
//! watermarker in [`ResampledWatermarker`] when the model expects a
//! different rate (e.g. AudioSeal trained at 16 kHz on a 24 kHz CSM stream).

use crate::audio_io;
use crate::error::Result;

pub trait Watermarker: Send + Sync {
    fn embed(&self, audio: &[f32]) -> Result<Vec<f32>>;

    /// Embed with a specific 16-bit message payload, overriding the
    /// watermarker's default. Default impl ignores `_message` and falls
    /// back to [`Self::embed`] — implementations that don't carry a
    /// payload (NoopWatermarker, future SilentCipher) inherit this no-op.
    /// Implementations that DO carry a payload (AudioSeal) override.
    fn embed_with_message(&self, audio: &[f32], _message: u16) -> Result<Vec<f32>> {
        self.embed(audio)
    }
}

#[derive(Default, Debug, Clone, Copy)]
pub struct NoopWatermarker;

impl Watermarker for NoopWatermarker {
    fn embed(&self, audio: &[f32]) -> Result<Vec<f32>> {
        Ok(audio.to_vec())
    }
}

/// Wraps a `Watermarker` that operates at `model_rate` so callers can use
/// it on audio at `source_rate`. Resamples in/out via the existing rubato
/// pipeline. Output length is normalized to the input length to keep this
/// as a drop-in equivalent of a same-rate watermarker.
pub struct ResampledWatermarker<W: Watermarker> {
    pub inner: W,
    pub source_rate: u32,
    pub model_rate: u32,
}

impl<W: Watermarker> ResampledWatermarker<W> {
    pub fn new(inner: W, source_rate: u32, model_rate: u32) -> Self {
        Self {
            inner,
            source_rate,
            model_rate,
        }
    }
}

impl<W: Watermarker> ResampledWatermarker<W> {
    fn run<F>(&self, audio: &[f32], inner_call: F) -> Result<Vec<f32>>
    where
        F: FnOnce(&[f32]) -> Result<Vec<f32>>,
    {
        if self.source_rate == self.model_rate {
            return inner_call(audio);
        }
        let n = audio.len();
        let down = audio_io::resample(audio, self.source_rate, self.model_rate)?;
        let watermarked = inner_call(&down)?;
        let mut up = audio_io::resample(&watermarked, self.model_rate, self.source_rate)?;
        if up.len() > n {
            up.truncate(n);
        } else if up.len() < n {
            up.resize(n, 0.0);
        }
        Ok(up)
    }
}

impl<W: Watermarker> Watermarker for ResampledWatermarker<W> {
    fn embed(&self, audio: &[f32]) -> Result<Vec<f32>> {
        self.run(audio, |a| self.inner.embed(a))
    }

    fn embed_with_message(&self, audio: &[f32], message: u16) -> Result<Vec<f32>> {
        self.run(audio, |a| self.inner.embed_with_message(a, message))
    }
}