//! Long-form generation: synthesize multi-sentence text by chunking, with a
//! rolling context that keeps the model coherent across chunks.
//!
//! Why this exists: CSM's max_seq_len = 2048 tokens (~2 min of audio history).
//! For utterances much longer than ~10–20 s the model also drifts in prosody
//! and voice characteristics. The mitigation is to split text on sentence
//! boundaries, generate one sentence at a time, and feed the previous
//! generated audio + transcript back as context for the next call.
//!
//! Anchor re-injection: every `anchor_every_n` chunks we re-prepend the
//! original speaker reference (if a `SpeakerProfile` is supplied) to combat
//! voice drift over long sequences.

use crate::error::Result;
use crate::generator::{GenerateOptions, Generator};
use crate::prompt::Segment;
use crate::speaker::SpeakerProfile;
use std::sync::OnceLock;

#[derive(Debug, Clone, Copy)]
pub struct LongFormConfig {
    /// Soft target for sentence-chunk character length.
    pub max_chunk_chars: usize,
    /// Drop oldest carry-back segments once total estimated tokens > this.
    pub rolling_context_budget: usize,
    /// Re-inject the speaker anchor every N chunks. 0 disables.
    pub anchor_every_n: usize,
}

impl Default for LongFormConfig {
    fn default() -> Self {
        Self {
            max_chunk_chars: 200,
            rolling_context_budget: 1500,
            anchor_every_n: 4,
        }
    }
}

/// Split `text` into sentence-shaped chunks bounded by `max_chunk_chars`.
pub fn split_sentences(text: &str, max_chunk_chars: usize) -> Vec<String> {
    let mut chunks: Vec<String> = Vec::new();
    let sentences = split_on_sentence_boundaries(text);
    let mut current = String::new();
    for s in sentences {
        let s = s.trim();
        if s.is_empty() {
            continue;
        }
        // If adding this sentence overflows AND we already have content, flush.
        if !current.is_empty() && current.chars().count() + 1 + s.chars().count() > max_chunk_chars
        {
            chunks.push(std::mem::take(&mut current));
        }
        if !current.is_empty() {
            current.push(' ');
        }
        current.push_str(s);
        // Long single sentence: hard-split at the cap to stay sane.
        while current.chars().count() > max_chunk_chars {
            let take = char_byte_index(&current, max_chunk_chars);
            // Try to backtrack to a space for a cleaner cut.
            let cut_at = current[..take].rfind(' ').unwrap_or(take);
            let head: String = current[..cut_at].into();
            let tail: String = current[cut_at..].trim_start().into();
            chunks.push(head);
            current = tail;
        }
    }
    if !current.is_empty() {
        chunks.push(current);
    }
    chunks
}

fn char_byte_index(s: &str, char_pos: usize) -> usize {
    s.char_indices()
        .nth(char_pos)
        .map(|(b, _)| b)
        .unwrap_or(s.len())
}

fn sentence_split_regex() -> &'static regex::Regex {
    static R: OnceLock<regex::Regex> = OnceLock::new();
    // Split after . ! ? followed by whitespace; keep the punctuation with the
    // preceding sentence by using a lookbehind-like trick (regex crate doesn't
    // support lookbehinds, so we capture and reattach).
    R.get_or_init(|| regex::Regex::new(r"(?P<end>[.!?]+)\s+").unwrap())
}

fn split_on_sentence_boundaries(text: &str) -> Vec<String> {
    let re = sentence_split_regex();
    let mut last = 0;
    let mut out: Vec<String> = Vec::new();
    for m in re.find_iter(text) {
        let end = m.end();
        out.push(text[last..end].to_string());
        last = end;
    }
    if last < text.len() {
        out.push(text[last..].to_string());
    }
    out
}

impl Generator {
    /// Generate audio for arbitrarily long text by chunking on sentence
    /// boundaries with a rolling context. The previous generated chunk's
    /// audio + transcript becomes context for the next call. Returns the
    /// full concatenated PCM.
    pub fn generate_long(
        &mut self,
        text: &str,
        speaker: u32,
        profile: Option<&SpeakerProfile>,
        opts: GenerateOptions,
        cfg: LongFormConfig,
    ) -> Result<Vec<f32>> {
        // We bypass the per-call text_normalize on the full text (it would
        // hard-cap at max_chars) and let each chunk normalize itself via the
        // generator's own pipeline.
        let normalized = self.text_normalize.apply(text)?;
        let chunks = split_sentences(&normalized, cfg.max_chunk_chars);
        if chunks.is_empty() {
            return Ok(Vec::new());
        }
        tracing::info!(
            "generate_long: {} chunks (avg {:.0} chars)",
            chunks.len(),
            normalized.chars().count() as f32 / chunks.len() as f32
        );

        let anchor_segments: Vec<Segment> =
            profile.map(|p| p.segments().to_vec()).unwrap_or_default();

        let mut rolling: Vec<Segment> = anchor_segments.clone();
        let mut full_pcm: Vec<f32> = Vec::new();

        for (chunk_idx, chunk_text) in chunks.iter().enumerate() {
            // Anchor re-injection.
            if cfg.anchor_every_n > 0
                && chunk_idx > 0
                && chunk_idx % cfg.anchor_every_n == 0
                && !anchor_segments.is_empty()
            {
                // Prepend anchors at the front of rolling context.
                rolling = {
                    let mut combined = anchor_segments.clone();
                    combined.extend(rolling.into_iter().filter(|s| {
                        // Avoid double-anchors if rolling already starts with anchor segments.
                        !anchor_segments
                            .iter()
                            .any(|a| std::ptr::eq(a as *const _, s as *const _))
                    }));
                    combined
                };
            }
            // Budget eviction of oldest non-anchor rolling context.
            evict_to_budget(&mut rolling, &anchor_segments, cfg.rolling_context_budget);

            tracing::info!(
                "  chunk {}/{}: {} chars, ctx={}",
                chunk_idx + 1,
                chunks.len(),
                chunk_text.chars().count(),
                rolling.len()
            );

            let pcm = self.generate(chunk_text, speaker, &rolling, opts.clone())?;
            full_pcm.extend_from_slice(&pcm);

            // Convert this chunk into a Segment and add to rolling context.
            let new_ctx = Segment::new(speaker, chunk_text.clone(), pcm);
            rolling.push(new_ctx);
        }

        Ok(full_pcm)
    }

    /// Long-form analogue of [`Self::generate_to_wav`]: chunked generation
    /// with rolling context, then post-processing, then optional watermark
    /// (if installed via [`Self::set_watermarker`]), then WAV write.
    ///
    /// Order matches `generate_to_wav` exactly so installing a watermarker
    /// applies uniformly to short-form and long-form output.
    #[allow(clippy::too_many_arguments)]
    pub fn generate_long_to_wav(
        &mut self,
        text: &str,
        speaker: u32,
        profile: Option<&SpeakerProfile>,
        opts: GenerateOptions,
        cfg: LongFormConfig,
        post: &crate::PostProcess,
        out_path: &std::path::Path,
    ) -> Result<()> {
        let mut pcm = self.generate_long(text, speaker, profile, opts, cfg)?;
        post.apply(&mut pcm, self.config.sample_rate)?;
        if let Some(wm) = self.watermarker.as_ref() {
            pcm = wm.embed(&pcm)?;
        }
        crate::audio_io::write_wav_24k_mono(out_path, &pcm)?;
        Ok(())
    }
}

fn estimate_segment_tokens(seg: &Segment) -> usize {
    let mut total = 0usize;
    if let Some(audio) = &seg.audio {
        total += audio.len().div_ceil(1920);
    }
    total += seg.text.chars().count().div_ceil(4);
    total
}

fn evict_to_budget(rolling: &mut Vec<Segment>, anchors: &[Segment], budget: usize) {
    // Anchors are immutable head; only evict from the post-anchor tail.
    let anchor_count = anchors.len();
    while rolling.len() > anchor_count + 1 {
        let total: usize = rolling.iter().map(estimate_segment_tokens).sum();
        if total <= budget {
            break;
        }
        // Drop oldest *non-anchor* segment.
        rolling.remove(anchor_count);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn split_simple_sentences() {
        let text = "First. Second! Third? Fourth.";
        let chunks = split_sentences(text, 200);
        assert_eq!(chunks.len(), 1, "all 4 fit in one 200-char chunk");
    }

    #[test]
    fn split_when_overflow() {
        let text = "First sentence here. Second sentence here. Third sentence here.";
        let chunks = split_sentences(text, 30);
        assert!(chunks.len() >= 2, "got {chunks:?}");
        for c in &chunks {
            assert!(c.chars().count() <= 35, "overflow: {c:?}");
        }
    }

    #[test]
    fn split_long_single_sentence_hard_breaks() {
        let text = "this is a very long sentence with no punctuation that just keeps going and going and going forever and ever and ever";
        let chunks = split_sentences(text, 30);
        assert!(chunks.len() >= 3, "got {chunks:?}");
        for c in &chunks {
            assert!(c.chars().count() <= 32, "overflow on {c:?}");
        }
    }

    #[test]
    fn split_empty_returns_empty() {
        let chunks = split_sentences("", 200);
        assert!(chunks.is_empty());
        let chunks = split_sentences("   \t\n  ", 200);
        assert!(chunks.is_empty());
    }

    #[test]
    fn evict_drops_oldest_post_anchor() {
        let anchors = vec![Segment::new_text(0, "anchor")];
        let mut rolling: Vec<Segment> = anchors.clone();
        // Simulate 3 large rolling segments
        for i in 0..3 {
            rolling.push(Segment::new(
                0,
                format!("chunk-{i}"),
                vec![0.0f32; 24_000 * 3],
            ));
        }
        let before = rolling.len();
        evict_to_budget(&mut rolling, &anchors, 50);
        assert!(rolling.len() < before);
        // First element is still the anchor.
        assert_eq!(rolling[0].text, "anchor");
    }
}