//! Medición de latencia por turno.
//!
//! El objetivo no es un histograma bonito sino contestar a una pregunta
//! concreta cada vez que el asistente responde: *¿quién se ha comido el
//! tiempo?* Por eso se marcan los cinco instantes que separan las etapas y se
//! imprime el desglose, en lugar de un único total que no dice dónde mirar.

use std::collections::BTreeMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use crate::event::TurnId;

/// Etapas del pipeline, en el orden en que ocurren.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Stage {
    /// Del final del habla a la transcripción definitiva.
    Asr,
    /// De la transcripción al primer fragmento de texto del modelo.
    LlmFirstToken,
    /// Del primer fragmento a la respuesta completa.
    LlmRest,
    /// Ejecución de herramientas.
    Tools,
    /// De la primera frase al primer audio recibido.
    TtsFirstAudio,
    /// Síntesis del resto del turno.
    TtsRest,
    /// Del final del habla al primer audio: lo que de verdad se percibe.
    PerceivedLatency,
}

impl Stage {
    pub fn label(self) -> &'static str {
        match self {
            Stage::Asr => "asr.final",
            Stage::LlmFirstToken => "llm.primer_token",
            Stage::LlmRest => "llm.resto",
            Stage::Tools => "herramientas",
            Stage::TtsFirstAudio => "tts.primer_audio",
            Stage::TtsRest => "tts.resto",
            Stage::PerceivedLatency => "LATENCIA PERCIBIDA",
        }
    }
}

/// Cronómetro de un turno. Se rellena a medida que el turno avanza y se
/// vuelca de una vez al terminar.
#[derive(Debug)]
pub struct TurnTimer {
    pub turn: TurnId,
    started: Instant,
    marks: BTreeMap<Stage, Duration>,
    /// Segundos de audio hablados por el usuario, para calcular el RTF del ASR.
    pub input_secs: f32,
    /// Segundos de audio sintetizados, para el RTF del TTS.
    pub output_secs: f32,
    pub sentences: usize,
    pub tool_calls: usize,
}

impl TurnTimer {
    pub fn new(turn: TurnId) -> Self {
        Self {
            turn,
            started: Instant::now(),
            marks: BTreeMap::new(),
            input_secs: 0.0,
            output_secs: 0.0,
            sentences: 0,
            tool_calls: 0,
        }
    }

    /// Instante de referencia del turno: el fin del habla del usuario.
    pub fn started(&self) -> Instant {
        self.started
    }

    pub fn record(&mut self, stage: Stage, took: Duration) {
        // Etapas que se repiten (una síntesis por frase) se acumulan; las que
        // marcan un hito (primer audio) se quedan con la primera medida, que
        // es la que describe la latencia de arranque.
        match stage {
            Stage::TtsFirstAudio | Stage::LlmFirstToken | Stage::PerceivedLatency => {
                self.marks.entry(stage).or_insert(took);
            }
            _ => *self.marks.entry(stage).or_default() += took,
        }
    }

    /// Marca una etapa con el tiempo transcurrido desde el inicio del turno.
    pub fn mark_since_start(&mut self, stage: Stage) {
        let took = self.started.elapsed();
        self.record(stage, took);
    }

    pub fn get(&self, stage: Stage) -> Option<Duration> {
        self.marks.get(&stage).copied()
    }

    pub fn total(&self) -> Duration {
        self.started.elapsed()
    }

    /// Etapa que más ha tardado, excluyendo el total percibido (que las agrega).
    pub fn bottleneck(&self) -> Option<(Stage, Duration)> {
        self.marks
            .iter()
            .filter(|(stage, _)| **stage != Stage::PerceivedLatency)
            .max_by_key(|(_, took)| **took)
            .map(|(stage, took)| (*stage, *took))
    }

    /// Informe de una línea por etapa, con el cuello de botella señalado.
    pub fn report(&self) -> String {
        let mut out = format!("turno {} — desglose de latencia\n", self.turn);
        let worst = self.bottleneck().map(|(stage, _)| stage);
        for (stage, took) in &self.marks {
            let flag = if Some(*stage) == worst {
                "  <== cuello de botella"
            } else {
                ""
            };
            out.push_str(&format!(
                "  {:<20} {:>8.0} ms{}\n",
                stage.label(),
                took.as_secs_f64() * 1000.0,
                flag
            ));
        }
        if self.input_secs > 0.0 {
            if let Some(asr) = self.get(Stage::Asr) {
                out.push_str(&format!(
                    "  {:<20} {:>8.2} x tiempo real ({:.1} s de voz)\n",
                    "asr.rtf",
                    asr.as_secs_f32() / self.input_secs,
                    self.input_secs
                ));
            }
        }
        if self.output_secs > 0.0 {
            let tts: Duration = self.get(Stage::TtsFirstAudio).unwrap_or_default()
                + self.get(Stage::TtsRest).unwrap_or_default();
            out.push_str(&format!(
                "  {:<20} {:>8.2} x tiempo real ({:.1} s de audio, {} frases)\n",
                "tts.rtf",
                tts.as_secs_f32() / self.output_secs,
                self.output_secs,
                self.sentences
            ));
        }
        out.push_str(&format!(
            "  {:<20} {:>8.0} ms\n",
            "total del turno",
            self.total().as_secs_f64() * 1000.0
        ));
        out
    }
}

/// Agregado de todos los turnos de la sesión, para el resumen del cierre.
#[derive(Debug, Default)]
pub struct Metrics {
    inner: Mutex<MetricsInner>,
}

#[derive(Debug, Default)]
struct MetricsInner {
    turns: usize,
    totals: BTreeMap<Stage, (Duration, usize)>,
}

impl Metrics {
    pub fn record_turn(&self, timer: &TurnTimer) {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.turns += 1;
        for (stage, took) in &timer.marks {
            let entry = inner.totals.entry(*stage).or_default();
            entry.0 += *took;
            entry.1 += 1;
        }
    }

    pub fn turns(&self) -> usize {
        self.inner.lock().unwrap_or_else(|e| e.into_inner()).turns
    }

    /// Media por etapa sobre toda la sesión.
    pub fn summary(&self) -> String {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if inner.turns == 0 {
            return "sin turnos completados\n".into();
        }
        let mut out = format!("resumen de {} turno(s) — media por etapa\n", inner.turns);
        for (stage, (total, count)) in &inner.totals {
            out.push_str(&format!(
                "  {:<20} {:>8.0} ms\n",
                stage.label(),
                total.as_secs_f64() * 1000.0 / *count as f64
            ));
        }
        out
    }
}

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

    #[test]
    fn el_cuello_de_botella_es_la_etapa_mas_lenta() {
        let mut timer = TurnTimer::new(TurnId(1));
        timer.record(Stage::Asr, Duration::from_millis(300));
        timer.record(Stage::LlmFirstToken, Duration::from_millis(500));
        timer.record(Stage::TtsFirstAudio, Duration::from_millis(900));
        assert_eq!(timer.bottleneck().unwrap().0, Stage::TtsFirstAudio);
    }

    #[test]
    fn la_latencia_percibida_no_compite_como_cuello_de_botella() {
        let mut timer = TurnTimer::new(TurnId(1));
        timer.record(Stage::Asr, Duration::from_millis(300));
        timer.record(Stage::PerceivedLatency, Duration::from_secs(9));
        assert_eq!(timer.bottleneck().unwrap().0, Stage::Asr);
    }

    #[test]
    fn las_etapas_repetidas_se_acumulan_y_los_hitos_no() {
        let mut timer = TurnTimer::new(TurnId(1));
        timer.record(Stage::TtsRest, Duration::from_millis(100));
        timer.record(Stage::TtsRest, Duration::from_millis(150));
        assert_eq!(timer.get(Stage::TtsRest), Some(Duration::from_millis(250)));

        timer.record(Stage::TtsFirstAudio, Duration::from_millis(600));
        timer.record(Stage::TtsFirstAudio, Duration::from_millis(50));
        assert_eq!(
            timer.get(Stage::TtsFirstAudio),
            Some(Duration::from_millis(600)),
            "el primer audio describe el arranque; una frase posterior no debe rebajarlo"
        );
    }

    #[test]
    fn el_resumen_promedia_entre_turnos() {
        let metrics = Metrics::default();
        for ms in [100, 300] {
            let mut timer = TurnTimer::new(TurnId(1));
            timer.record(Stage::Asr, Duration::from_millis(ms));
            metrics.record_turn(&timer);
        }
        assert_eq!(metrics.turns(), 2);
        assert!(metrics.summary().contains("200 ms"));
    }
}