//! Composable orchestrator for the LLM → TTS half of the conversational
//! stack. Takes a prompt + chat history, streams LLM tokens, buffers them
//! into sentences, and dispatches each completed sentence to CSM TTS as
//! soon as it's ready. PCM is emitted as the model produces it.
//!
//! ## Why sentence buffering?
//!
//! CSM is a sentence-level TTS — its prosody is best when fed a complete
//! sentence at a time, not token-by-token. The LLM streams tokens; we
//! accumulate until a terminal punctuation mark (`.`, `!`, `?`) or newline,
//! then flush the sentence to TTS. This trades a small chunk of buffering
//! latency (typically <1s for short sentences) for natural prosody.
//!
//! ## Pipeline
//!
//! ```text
//!   prompt + history
//!     │
//!     ▼
//!   LlmClient.generate_stream → token chunks (k1, k2, ...)
//!     │
//!     ▼ buffer until terminal punctuation
//!   Sentence "Hello."
//!     │
//!     ▼
//!   Generator.generate(text) → 24 kHz PCM (Vec<f32>)
//!     │
//!     ▼ on_audio callback fires once per sentence
//!   caller (e.g. WebSocket sink, file writer)
//! ```
//!
//! When the next implementation lands a streaming TTS variant that takes
//! token-level input directly (no sentence boundary needed), the
//! orchestrator can be swapped to that path. Today's CSM is the bottleneck.
//!
//! ## Half-duplex; full duplex (6c.2) is deferred
//!
//! This orchestrator handles the LLM → TTS direction only. To get full
//! voice-in/voice-out, pair this with the Phase 6a STT once its word
//! emission is fixed: `audio_in → STT → user_text → Converse → audio_out`.

use crate::GenerateOptions;
use crate::audio_io;
use crate::error::{CsmError, Result};
use crate::generator::Generator;
use crate::llm_client::{ChatMessage, GenConfig, LlmClient};
use crate::post::PostProcess;
use futures_util::StreamExt;
use std::path::Path;

/// One unit emitted by [`Converse::run`] each time a sentence completes.
#[derive(Debug, Clone)]
pub struct Utterance {
    pub text: String,
    /// 24 kHz mono PCM, post-processed and (if a watermarker is installed
    /// on the Generator) watermarked.
    pub audio: Vec<f32>,
    /// Wall-clock time-to-first-audio for THIS sentence, in milliseconds.
    /// (Time from sentence-buffer flush to PCM produced.)
    pub tts_latency_ms: u128,
}

/// How aggressively to flush sentences. `Punctuation` waits for `. ! ?`
/// or newline; `Eager` flushes more often for lower latency at the cost
/// of less natural prosody.
#[derive(Debug, Clone, Copy)]
pub enum FlushPolicy {
    Punctuation,
    /// Flush at every comma, semicolon, or punctuation mark — useful for
    /// long-running monologues where you want low first-audio latency.
    Eager,
}

#[derive(Debug, Clone)]
pub struct ConverseOptions {
    pub generate: GenerateOptions,
    pub speaker: u32,
    pub flush: FlushPolicy,
    /// Skip sentences shorter than this (after trim). Avoids dispatching
    /// tiny "k." or single-character fragments to TTS.
    pub min_sentence_chars: usize,
}

impl Default for ConverseOptions {
    fn default() -> Self {
        Self {
            generate: GenerateOptions {
                max_audio_ms: 6_000,
                ..GenerateOptions::default()
            },
            speaker: 0,
            flush: FlushPolicy::Punctuation,
            min_sentence_chars: 2,
        }
    }
}

/// Locate the byte index AFTER the first sentence-boundary character in
/// the buffer per policy. Returns `Some(end)` so the caller can do
/// `buf.drain(..end)` to extract the completed sentence (including the
/// punctuation mark and any trailing whitespace up to the boundary).
/// Returns `None` if no boundary is present yet.
fn find_first_boundary(buf: &str, policy: FlushPolicy) -> Option<usize> {
    let is_boundary = |c: char| match policy {
        FlushPolicy::Punctuation => matches!(c, '.' | '!' | '?' | '\n'),
        FlushPolicy::Eager => matches!(c, '.' | '!' | '?' | '\n' | ',' | ';' | ':'),
    };
    for (i, c) in buf.char_indices() {
        if is_boundary(c) {
            return Some(i + c.len_utf8());
        }
    }
    None
}

/// Buffer contains a flushable boundary anywhere (used by the
/// streaming-path tests). For the actual flush logic we use
/// [`find_first_boundary`] to get the byte index of the first boundary.
#[cfg(test)]
fn should_flush(buf: &str, policy: FlushPolicy) -> bool {
    find_first_boundary(buf, policy).is_some()
}

/// Optional hook invoked with `&mut Generator` before each sentence's
/// synthesis call. Use this to apply per-sentence steering (e.g. emotion
/// shifts mid-reply). Sync — apply_steering is non-async — so the hook
/// fits cleanly between LLM token consumption and the synthesize call.
pub type PreSentenceHook = Box<dyn FnMut(&mut Generator, &str) -> Result<()> + Send + 'static>;

pub struct Converse<'a, L: LlmClient> {
    llm: &'a L,
    generator: &'a mut Generator,
    post: PostProcess,
    /// Persistent speaker-prompt context prepended to every sentence's
    /// generation call. Without this, CSM-1B has no voice anchor and
    /// drifts between speakers / pitches across turns. Typically a
    /// single `Segment` carrying a 5–15s reference clip with its
    /// transcript pins the voice for the whole session.
    context: Vec<crate::prompt::Segment>,
    /// Hook fired before each sentence's `synthesize_streaming` call.
    /// Receives mutable access to the underlying Generator and the
    /// sentence text, so callers can apply per-sentence steering or
    /// other state changes without touching this crate's internals.
    pre_sentence_hook: Option<PreSentenceHook>,
}

impl<'a, L: LlmClient> Converse<'a, L> {
    pub fn new(llm: &'a L, generator: &'a mut Generator) -> Self {
        Self {
            llm,
            generator,
            post: PostProcess::default(),
            context: Vec::new(),
            pre_sentence_hook: None,
        }
    }

    /// Install a hook that fires before each sentence's synth call,
    /// with mutable access to the underlying Generator + the sentence
    /// text. Use this to apply per-sentence steering (e.g. emotion
    /// shifts mid-reply).
    pub fn with_pre_sentence_hook(mut self, hook: PreSentenceHook) -> Self {
        self.pre_sentence_hook = Some(hook);
        self
    }

    /// Override post-processing (HPF + declick + LUFS). Pass
    /// [`PostProcess::disabled`] to skip.
    pub fn with_post(mut self, post: PostProcess) -> Self {
        self.post = post;
        self
    }

    /// Set persistent speaker-prompt context for the session.
    /// Each segment provides a (speaker, text, audio) triple the model
    /// uses as a voice anchor. Without this, CSM-1B picks a different
    /// speaker per turn.
    pub fn with_context(mut self, ctx: Vec<crate::prompt::Segment>) -> Self {
        self.context = ctx;
        self
    }

    /// Run the full pipeline: stream LLM tokens, flush sentences to TTS,
    /// invoke `on_utterance` for each completed (text, audio) pair as
    /// they're produced. Returns the full assistant message text once
    /// the LLM stream ends.
    pub async fn run<F>(
        &mut self,
        messages: Vec<ChatMessage>,
        gen_cfg: GenConfig,
        opts: ConverseOptions,
        mut on_utterance: F,
    ) -> Result<String>
    where
        F: FnMut(&Utterance) -> Result<()>,
    {
        // Phase tracing: TTFT (time-to-first-token) for the LLM, and per-
        // sentence llm_buffer (time accumulating tokens until a boundary)
        // and tts_gen (time inside synthesize). All emitted at info level
        // with a stable "conv-phase:" prefix so callers can filter or grep.
        let run_t = std::time::Instant::now();
        let mut stream = self.llm.generate_stream(messages, gen_cfg).await?;
        let mut buf = String::new();
        let mut full = String::new();
        let mut first_token_logged = false;
        let mut sentence_idx: usize = 0;
        let mut sentence_t = std::time::Instant::now();

        while let Some(chunk) = stream.next().await {
            let chunk = chunk?;
            if !first_token_logged {
                first_token_logged = true;
                tracing::info!(
                    "conv-phase: ttft={}ms (LLM first token)",
                    run_t.elapsed().as_millis()
                );
                sentence_t = std::time::Instant::now();
            }
            full.push_str(&chunk);
            buf.push_str(&chunk);
            // Flush as many complete sentences as the buffer contains.
            while let Some(end) = find_first_boundary(&buf, opts.flush) {
                let sentence: String = buf.drain(..end).collect();
                let trimmed = sentence.trim();
                if trimmed.chars().count() < opts.min_sentence_chars {
                    continue;
                }
                let llm_buffer_ms = sentence_t.elapsed().as_millis();
                if let Some(hook) = self.pre_sentence_hook.as_mut() {
                    hook(self.generator, trimmed)?;
                }
                let utt = self.synthesize(trimmed, &opts)?;
                tracing::info!(
                    "conv-phase: sentence[{sentence_idx}] llm_buffer={llm_buffer_ms}ms \
                     tts_gen={}ms chars={}",
                    utt.tts_latency_ms,
                    trimmed.chars().count(),
                );
                sentence_idx += 1;
                sentence_t = std::time::Instant::now();
                on_utterance(&utt)?;
            }
        }

        // Trailing text without terminal punctuation — flush as one
        // final sentence so callers don't lose the tail.
        let trailing = buf.trim().to_string();
        if trailing.chars().count() >= opts.min_sentence_chars {
            let llm_buffer_ms = sentence_t.elapsed().as_millis();
            if let Some(hook) = self.pre_sentence_hook.as_mut() {
                hook(self.generator, &trailing)?;
            }
            let utt = self.synthesize(&trailing, &opts)?;
            tracing::info!(
                "conv-phase: sentence[{sentence_idx}] (trailing) llm_buffer={llm_buffer_ms}ms \
                 tts_gen={}ms chars={}",
                utt.tts_latency_ms,
                trailing.chars().count(),
            );
            on_utterance(&utt)?;
        }
        Ok(full)
    }

    fn synthesize(&mut self, sentence: &str, opts: &ConverseOptions) -> Result<Utterance> {
        let t = std::time::Instant::now();
        let mut pcm = self
            .generator
            .generate(sentence, opts.speaker, &self.context, opts.generate.clone())
            .map_err(|e| CsmError::Config(format!("converse generate: {e}")))?;
        self.post
            .apply(&mut pcm, self.generator.config.sample_rate)
            .map_err(|e| CsmError::Config(format!("converse post: {e}")))?;
        if let Some(wm) = self.generator.watermarker.as_ref() {
            pcm = wm.embed(&pcm)?;
        }
        Ok(Utterance {
            text: sentence.to_string(),
            audio: pcm,
            tts_latency_ms: t.elapsed().as_millis(),
        })
    }

    /// Streaming variant of [`Self::run`]: for each sentence, the
    /// underlying TTS uses [`Generator::generate_streaming`] with
    /// `chunk_frames`-sized chunks (4 frames ≈ 320 ms at 12.5 Hz).
    /// `on_chunk` fires for each chunk as it's produced — this is the
    /// path that delivers low first-audio latency. `on_sentence` fires
    /// once per completed sentence with the accumulated audio (useful
    /// for /metrics + transcript bookkeeping; the audio has already
    /// been streamed out via `on_chunk`).
    ///
    /// **No post-process or watermark in this mode.** Both operations
    /// require full-sentence context (LUFS normalize over the whole
    /// utterance, AudioSeal needs ~1s of audio for a stable embed) and
    /// can't run on partial chunks without re-sending audio that has
    /// already left the server. Callers needing those features should
    /// use [`Self::run`] instead. If a watermarker is installed on the
    /// underlying generator, this method returns an error to make the
    /// tradeoff explicit.
    pub async fn run_streaming<C, U>(
        &mut self,
        messages: Vec<ChatMessage>,
        gen_cfg: GenConfig,
        opts: ConverseOptions,
        chunk_frames: usize,
        mut on_chunk: C,
        mut on_sentence: U,
    ) -> Result<String>
    where
        C: FnMut(&[f32]) -> Result<()>,
        U: FnMut(&Utterance) -> Result<()>,
    {
        if self.generator.watermarker.is_some() {
            return Err(CsmError::Config(
                "run_streaming: watermarker installed but streaming mode \
                 cannot embed (needs full-sentence context). Use run() \
                 instead, or remove the watermarker."
                    .into(),
            ));
        }
        let run_t = std::time::Instant::now();
        let mut stream = self.llm.generate_stream(messages, gen_cfg).await?;
        let mut buf = String::new();
        let mut full = String::new();
        let mut first_token_logged = false;
        let mut sentence_idx: usize = 0;
        let mut sentence_t = std::time::Instant::now();

        while let Some(chunk) = stream.next().await {
            let chunk = chunk?;
            if !first_token_logged {
                first_token_logged = true;
                tracing::info!(
                    "conv-phase: ttft={}ms (LLM first token, streaming)",
                    run_t.elapsed().as_millis()
                );
                sentence_t = std::time::Instant::now();
            }
            full.push_str(&chunk);
            buf.push_str(&chunk);
            while let Some(end) = find_first_boundary(&buf, opts.flush) {
                let sentence: String = buf.drain(..end).collect();
                let trimmed = sentence.trim();
                if trimmed.chars().count() < opts.min_sentence_chars {
                    continue;
                }
                let llm_buffer_ms = sentence_t.elapsed().as_millis();
                if let Some(hook) = self.pre_sentence_hook.as_mut() {
                    hook(self.generator, trimmed)?;
                }
                let utt = self.synthesize_streaming(trimmed, &opts, chunk_frames, &mut on_chunk)?;
                tracing::info!(
                    "conv-phase: sentence[{sentence_idx}] (streaming) \
                     llm_buffer={llm_buffer_ms}ms tts_gen={}ms chars={}",
                    utt.tts_latency_ms,
                    trimmed.chars().count(),
                );
                sentence_idx += 1;
                sentence_t = std::time::Instant::now();
                on_sentence(&utt)?;
            }
        }
        let trailing = buf.trim().to_string();
        if trailing.chars().count() >= opts.min_sentence_chars {
            let llm_buffer_ms = sentence_t.elapsed().as_millis();
            if let Some(hook) = self.pre_sentence_hook.as_mut() {
                hook(self.generator, &trailing)?;
            }
            let utt = self.synthesize_streaming(&trailing, &opts, chunk_frames, &mut on_chunk)?;
            tracing::info!(
                "conv-phase: sentence[{sentence_idx}] (streaming, trailing) \
                 llm_buffer={llm_buffer_ms}ms tts_gen={}ms chars={}",
                utt.tts_latency_ms,
                trailing.chars().count(),
            );
            on_sentence(&utt)?;
        }
        Ok(full)
    }

    fn synthesize_streaming<C>(
        &mut self,
        sentence: &str,
        opts: &ConverseOptions,
        chunk_frames: usize,
        on_chunk: &mut C,
    ) -> Result<Utterance>
    where
        C: FnMut(&[f32]) -> Result<()>,
    {
        let t = std::time::Instant::now();
        // generate_streaming is synchronous. The chunk callback fires
        // per-chunk; chunks are forwarded via whatever channel the
        // caller wired up. For chunks to actually flush to the wire as
        // they're decoded (rather than batched at sentence boundary),
        // the caller must arrange for the *consumer* to live on a
        // different task than this `run_streaming` future — e.g., spawn
        // it via `tokio::spawn` rather than `tokio::join!`. Without
        // that, chunks pile up in the channel until this sync block
        // returns control to the runtime.
        let pcm = self
            .generator
            .generate_streaming(
                sentence,
                opts.speaker,
                &self.context,
                opts.generate.clone(),
                chunk_frames,
                |chunk| on_chunk(chunk),
            )
            .map_err(|e| CsmError::Config(format!("converse generate_streaming: {e}")))?;
        Ok(Utterance {
            text: sentence.to_string(),
            audio: pcm,
            tts_latency_ms: t.elapsed().as_millis(),
        })
    }
}

/// Convenience: write all utterances concatenated into a single 24 kHz WAV.
pub fn write_concatenated_wav<I>(utterances: I, out: &Path) -> Result<()>
where
    I: IntoIterator<Item = Utterance>,
{
    let mut all: Vec<f32> = Vec::new();
    for u in utterances {
        all.extend(u.audio);
    }
    audio_io::write_wav_24k_mono(out, &all)
}

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

    #[test]
    fn flush_punctuation_modes() {
        assert!(should_flush("Hello.", FlushPolicy::Punctuation));
        assert!(should_flush("Hello. World", FlushPolicy::Punctuation));
        assert!(should_flush("Wait!", FlushPolicy::Punctuation));
        assert!(should_flush("What?", FlushPolicy::Punctuation));
        assert!(should_flush("Line\n", FlushPolicy::Punctuation));
        assert!(!should_flush("Hello", FlushPolicy::Punctuation));
        assert!(!should_flush("Hello, world", FlushPolicy::Punctuation));
        assert!(should_flush("Hello,", FlushPolicy::Eager));
        assert!(should_flush("Hello, world", FlushPolicy::Eager));
        assert!(should_flush("then;", FlushPolicy::Eager));
        assert!(!should_flush("Hello world", FlushPolicy::Eager));
    }

    #[test]
    fn flush_empty_buffer() {
        assert!(!should_flush("", FlushPolicy::Punctuation));
        assert!(!should_flush("", FlushPolicy::Eager));
    }

    #[test]
    fn find_boundary_returns_index_after_punctuation() {
        // "Hello. World" → boundary at index 6 (after the .).
        assert_eq!(
            find_first_boundary("Hello. World", FlushPolicy::Punctuation),
            Some(6)
        );
        assert_eq!(find_first_boundary("Hello", FlushPolicy::Punctuation), None);
        // Eager: comma at index 5, returns 6.
        assert_eq!(
            find_first_boundary("Hello, world", FlushPolicy::Eager),
            Some(6)
        );
    }
}