//! Text-to-speech engine — synthesizes text to WAV files using Kokoro neural TTS.
//!
//! The model is loaded lazily on first use. If the model/voices files are not
//! present, the generate command returns an error and the caller skips the prompt.
//!
//! Caching is handled internally via a `.meta` sidecar file next to each WAV.
//! When `cacheable` is true, the engine checks whether the existing WAV was
//! generated from the same text+voice; if so it returns immediately (cache hit).
//! Callers never need to check for cached files — that is entirely this module's
//! responsibility.

use crate::audio_player::pcm_to_mix_frames;
use kokoro_tts::{KokoroTts, Voice};
use std::path::Path;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::{mpsc, watch};

pub const DEFAULT_MODEL_PATH: &str = ".nogit/tts/kokoro-v1.0.onnx";
pub const DEFAULT_VOICES_PATH: &str = ".nogit/tts/voices.bin";
const TTS_OUTPUT_RATE: u32 = 24000;
const MAX_CHUNK_CHARS: usize = 220;
const MIN_CHUNK_CHARS: usize = 80;

pub enum TtsStreamMessage {
    Frames(Vec<Vec<f32>>),
    Finished,
    Failed(String),
}

pub struct TtsLivePrompt {
    pub initial_frames: Vec<Vec<f32>>,
    pub stream_rx: mpsc::Receiver<TtsStreamMessage>,
    pub cancel_tx: watch::Sender<bool>,
}

#[derive(Clone)]
pub struct TtsPromptRequest {
    pub model_path: String,
    pub voices_path: String,
    pub voice_name: String,
    pub text: String,
}

/// Wraps the Kokoro TTS engine with lazy model loading.
pub struct TtsEngine {
    tts: Option<Arc<KokoroTts>>,
    /// Path that was used to load the current model (for cache invalidation).
    loaded_model_path: String,
    loaded_voices_path: String,
    /// On-disk TTS WAVs are cacheable only within a single engine lifetime.
    /// Every restart gets a new generation token, so prior process outputs are
    /// treated as stale and regenerated on first use.
    cache_generation: String,
}

impl TtsEngine {
    pub fn new() -> Self {
        Self {
            tts: None,
            loaded_model_path: String::new(),
            loaded_voices_path: String::new(),
            cache_generation: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .map(|d| d.as_nanos().to_string())
                .unwrap_or_else(|_| "0".to_string()),
        }
    }

    async fn ensure_loaded(
        &mut self,
        model_path: &str,
        voices_path: &str,
    ) -> Result<Arc<KokoroTts>, String> {
        if !Path::new(model_path).exists() {
            return Err(format!("model not found: {model_path}"));
        }
        if !Path::new(voices_path).exists() {
            return Err(format!("voices not found: {voices_path}"));
        }

        if self.tts.is_none()
            || self.loaded_model_path != model_path
            || self.loaded_voices_path != voices_path
        {
            eprintln!("[tts] loading model: {model_path}");
            let tts = Arc::new(
                KokoroTts::new(model_path, voices_path)
                    .await
                    .map_err(|e| format!("model load failed: {e:?}"))?,
            );
            self.tts = Some(tts);
            self.loaded_model_path = model_path.to_string();
            self.loaded_voices_path = voices_path.to_string();
        }

        Ok(self.tts.as_ref().unwrap().clone())
    }

    pub async fn start_live_prompt(
        &mut self,
        request: TtsPromptRequest,
    ) -> Result<TtsLivePrompt, String> {
        if request.text.trim().is_empty() {
            return Err("empty text".into());
        }

        let tts = self
            .ensure_loaded(&request.model_path, &request.voices_path)
            .await?;
        let voice = select_voice(&request.voice_name);
        let chunks = chunk_text(&request.text);
        if chunks.is_empty() {
            return Err("empty text".into());
        }

        let initial_frames = synth_text_to_mix_frames(&tts, chunks[0].as_str(), voice).await?;
        let remaining_chunks: Vec<String> = chunks.into_iter().skip(1).collect();
        let (stream_tx, stream_rx) = mpsc::channel(8);
        let (cancel_tx, cancel_rx) = watch::channel(false);

        tokio::spawn(async move {
            stream_live_prompt_chunks(tts, voice, remaining_chunks, stream_tx, cancel_rx).await;
        });

        Ok(TtsLivePrompt {
            initial_frames,
            stream_rx,
            cancel_tx,
        })
    }

    /// Generate a WAV file from text.
    ///
    /// Params (from IPC JSON):
    ///   - `model`: path to the ONNX model file
    ///   - `voices`: path to the voices.bin file
    ///   - `voice`: voice name (e.g. "af_bella")
    ///   - `text`: text to synthesize
    ///   - `output`: output WAV file path
    ///   - `cacheable`: if true, skip synthesis when the output WAV already
    ///     matches the same text+voice (checked via a `.meta` sidecar file)
    pub async fn generate(
        &mut self,
        params: &serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        let model_path = params
            .get("model")
            .and_then(|v| v.as_str())
            .ok_or("missing 'model' param")?;
        let voices_path = params
            .get("voices")
            .and_then(|v| v.as_str())
            .ok_or("missing 'voices' param")?;
        let voice_name = params
            .get("voice")
            .and_then(|v| v.as_str())
            .unwrap_or("af_bella");
        let text = params
            .get("text")
            .and_then(|v| v.as_str())
            .ok_or("missing 'text' param")?;
        let output_path = params
            .get("output")
            .and_then(|v| v.as_str())
            .ok_or("missing 'output' param")?;
        let cacheable = params
            .get("cacheable")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        if text.is_empty() {
            return Err("empty text".into());
        }

        // Cache check: if cacheable and the sidecar matches, return immediately.
        if cacheable && self.is_cache_hit(output_path, text, voice_name) {
            eprintln!("[tts] cache hit: {output_path}");
            return Ok(serde_json::json!({ "output": output_path }));
        }

        // Ensure parent directory exists.
        if let Some(parent) = Path::new(output_path).parent() {
            let _ = std::fs::create_dir_all(parent);
        }

        let tts = self.ensure_loaded(model_path, voices_path).await?;
        let voice = select_voice(voice_name);

        eprintln!("[tts] synthesizing WAV voice '{voice_name}' to {output_path}");
        let (samples, duration) = tts
            .synth(text, voice)
            .await
            .map_err(|e| format!("synthesis failed: {e:?}"))?;
        eprintln!(
            "[tts] synthesized {} samples in {duration:?}",
            samples.len()
        );

        // Write 24kHz 16-bit mono WAV.
        let spec = hound::WavSpec {
            channels: 1,
            sample_rate: 24000,
            bits_per_sample: 16,
            sample_format: hound::SampleFormat::Int,
        };

        let mut writer = hound::WavWriter::create(output_path, spec)
            .map_err(|e| format!("WAV create failed: {e}"))?;
        for &sample in &samples {
            let s16 = (sample * 32767.0).round().clamp(-32768.0, 32767.0) as i16;
            writer
                .write_sample(s16)
                .map_err(|e| format!("WAV write: {e}"))?;
        }
        writer
            .finalize()
            .map_err(|e| format!("WAV finalize: {e}"))?;

        // Write sidecar for future cache checks.
        if cacheable {
            self.write_cache_meta(output_path, text, voice_name);
        }

        eprintln!("[tts] wrote {output_path}");
        Ok(serde_json::json!({ "output": output_path }))
    }

    // -----------------------------------------------------------------------
    // Cache helpers
    // -----------------------------------------------------------------------

    /// Check if the WAV + sidecar on disk match the given text+voice.
    fn is_cache_hit(&self, output_path: &str, text: &str, voice: &str) -> bool {
        let meta_path = format!("{output_path}.meta");
        if !Path::new(output_path).exists() || !Path::new(&meta_path).exists() {
            return false;
        }
        match std::fs::read_to_string(&meta_path) {
            Ok(contents) => contents == self.cache_key(text, voice),
            Err(_) => false,
        }
    }

    /// Write the sidecar `.meta` file next to the WAV.
    fn write_cache_meta(&self, output_path: &str, text: &str, voice: &str) {
        let meta_path = format!("{output_path}.meta");
        let _ = std::fs::write(&meta_path, self.cache_key(text, voice));
    }

    /// Build the cache key from process generation + text + voice.
    fn cache_key(&self, text: &str, voice: &str) -> String {
        format!("{}\0{}\0{}", self.cache_generation, text, voice)
    }
}

async fn synth_text_to_mix_frames(
    tts: &Arc<KokoroTts>,
    text: &str,
    voice: Voice,
) -> Result<Vec<Vec<f32>>, String> {
    let (samples, duration) = tts
        .synth(text, voice)
        .await
        .map_err(|e| format!("synthesis failed: {e:?}"))?;
    eprintln!(
        "[tts] synthesized chunk ({} chars, {} samples) in {duration:?}",
        text.chars().count(),
        samples.len()
    );
    pcm_to_mix_frames(&samples, TTS_OUTPUT_RATE)
}

async fn stream_live_prompt_chunks(
    tts: Arc<KokoroTts>,
    voice: Voice,
    chunks: Vec<String>,
    stream_tx: mpsc::Sender<TtsStreamMessage>,
    mut cancel_rx: watch::Receiver<bool>,
) {
    for chunk in chunks {
        if *cancel_rx.borrow() {
            break;
        }

        match synth_text_to_mix_frames(&tts, &chunk, voice).await {
            Ok(frames) => {
                if *cancel_rx.borrow() {
                    break;
                }
                if stream_tx.send(TtsStreamMessage::Frames(frames)).await.is_err() {
                    return;
                }
            }
            Err(error) => {
                let _ = stream_tx.send(TtsStreamMessage::Failed(error)).await;
                return;
            }
        }

        if cancel_rx.has_changed().unwrap_or(false) && *cancel_rx.borrow_and_update() {
            break;
        }
    }

    let _ = stream_tx.send(TtsStreamMessage::Finished).await;
}

fn chunk_text(text: &str) -> Vec<String> {
    let mut chunks = Vec::new();
    let mut current = String::new();

    for ch in text.chars() {
        current.push(ch);

        let len = current.chars().count();
        let hard_split = len >= MAX_CHUNK_CHARS && (ch.is_whitespace() || is_soft_boundary(ch));
        let natural_split = len >= MIN_CHUNK_CHARS && is_sentence_boundary(ch);

        if natural_split || hard_split {
            push_chunk(&mut chunks, &mut current);
        }
    }

    push_chunk(&mut chunks, &mut current);

    if chunks.len() >= 2 {
        let last_len = chunks.last().unwrap().chars().count();
        if last_len < (MIN_CHUNK_CHARS / 2) {
            let tail = chunks.pop().unwrap();
            if let Some(prev) = chunks.last_mut() {
                prev.push(' ');
                prev.push_str(tail.trim());
            }
        }
    }

    chunks
}

fn push_chunk(chunks: &mut Vec<String>, current: &mut String) {
    let trimmed = current.trim();
    if !trimmed.is_empty() {
        chunks.push(trimmed.to_string());
    }
    current.clear();
}

fn is_sentence_boundary(ch: char) -> bool {
    matches!(ch, '.' | '!' | '?' | '\n' | ';' | ':')
}

fn is_soft_boundary(ch: char) -> bool {
    matches!(ch, ',' | ';' | ':' | ')' | ']' | '\n')
}

/// Map voice name string to Kokoro Voice enum variant.
fn select_voice(name: &str) -> Voice {
    match name {
        "af_bella" => Voice::AfBella(1.0),
        "af_heart" => Voice::AfHeart(1.0),
        "af_jessica" => Voice::AfJessica(1.0),
        "af_nicole" => Voice::AfNicole(1.0),
        "af_nova" => Voice::AfNova(1.0),
        "af_sarah" => Voice::AfSarah(1.0),
        "af_sky" => Voice::AfSky(1.0),
        "af_river" => Voice::AfRiver(1.0),
        "af_alloy" => Voice::AfAlloy(1.0),
        "af_aoede" => Voice::AfAoede(1.0),
        "af_kore" => Voice::AfKore(1.0),
        "am_adam" => Voice::AmAdam(1.0),
        "am_echo" => Voice::AmEcho(1.0),
        "am_eric" => Voice::AmEric(1.0),
        "am_fenrir" => Voice::AmFenrir(1.0),
        "am_liam" => Voice::AmLiam(1.0),
        "am_michael" => Voice::AmMichael(1.0),
        "am_onyx" => Voice::AmOnyx(1.0),
        "am_puck" => Voice::AmPuck(1.0),
        "bf_alice" => Voice::BfAlice(1.0),
        "bf_emma" => Voice::BfEmma(1.0),
        "bf_isabella" => Voice::BfIsabella(1.0),
        "bf_lily" => Voice::BfLily(1.0),
        "bm_daniel" => Voice::BmDaniel(1.0),
        "bm_fable" => Voice::BmFable(1.0),
        "bm_george" => Voice::BmGeorge(1.0),
        "bm_lewis" => Voice::BmLewis(1.0),
        _ => {
            eprintln!("[tts] unknown voice '{name}', falling back to af_bella");
            Voice::AfBella(1.0)
        }
    }
}
