//! Per-turn latency measurement.
//!
//! The goal is not a pretty histogram but answering one concrete question every
//! time the assistant replies: *who ate the time?* That is why the five
//! instants separating the stages are marked and the breakdown is printed,
//! instead of a single total that does not say where to look.

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

use crate::event::TurnId;

/// Pipeline stages, in the order they happen.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Stage {
    /// From end of speech to the final transcription.
    Asr,
    /// From the transcription to the model's first text fragment.
    LlmFirstToken,
    /// From the first fragment to the complete answer.
    LlmRest,
    /// Tool execution.
    Tools,
    /// From the first sentence to the first audio received.
    TtsFirstAudio,
    /// Synthesis of the rest of the turn.
    TtsRest,
    /// From end of speech to first audio: what is really perceived.
    PerceivedLatency,
}

impl Stage {
    pub fn label(self) -> &'static str {
        match self {
            Stage::Asr => "asr.final",
            Stage::LlmFirstToken => "llm.first_token",
            Stage::LlmRest => "llm.rest",
            Stage::Tools => "tools",
            Stage::TtsFirstAudio => "tts.first_audio",
            Stage::TtsRest => "tts.rest",
            Stage::PerceivedLatency => "PERCEIVED LATENCY",
        }
    }
}

/// Stopwatch for one turn. Filled in as the turn progresses and dumped
/// in one go at the end.
#[derive(Debug)]
pub struct TurnTimer {
    pub turn: TurnId,
    started: Instant,
    marks: BTreeMap<Stage, Duration>,
    /// Seconds of audio spoken by the user, for the ASR RTF.
    pub input_secs: f32,
    /// Seconds of audio synthesized, for the TTS RTF.
    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,
        }
    }

    /// Reference instant of the turn: the end of the user's speech.
    pub fn started(&self) -> Instant {
        self.started
    }

    pub fn record(&mut self, stage: Stage, took: Duration) {
        // Stages that repeat (one synthesis per sentence) accumulate; the ones
        // that mark a milestone (first audio) keep the first measurement, which
        // is the one describing startup latency.
        match stage {
            Stage::TtsFirstAudio | Stage::LlmFirstToken | Stage::PerceivedLatency => {
                self.marks.entry(stage).or_insert(took);
            }
            _ => *self.marks.entry(stage).or_default() += took,
        }
    }

    /// Marks a stage with the time elapsed since the start of the turn.
    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()
    }

    /// Slowest stage, excluding the perceived total (which aggregates them).
    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))
    }

    /// One-line-per-stage report, with the bottleneck flagged.
    pub fn report(&self) -> String {
        let mut out = format!("turn {} — latency breakdown\n", self.turn);
        let worst = self.bottleneck().map(|(stage, _)| stage);
        for (stage, took) in &self.marks {
            let flag = if Some(*stage) == worst {
                "  <== bottleneck"
            } 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 real time ({:.1} s of speech)\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 real time ({:.1} s of audio, {} sentences)\n",
                "tts.rtf",
                tts.as_secs_f32() / self.output_secs,
                self.output_secs,
                self.sentences
            ));
        }
        out.push_str(&format!(
            "  {:<20} {:>8.0} ms\n",
            "turn total",
            self.total().as_secs_f64() * 1000.0
        ));
        out
    }
}

/// Aggregate of every turn in the session, for the closing summary.
#[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
    }

    /// Per-stage average over the whole session.
    pub fn summary(&self) -> String {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if inner.turns == 0 {
            return "no completed turns\n".into();
        }
        let mut out = format!("summary of {} turn(s) — average per stage\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 bottleneck_is_the_slowest_stage() {
        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 perceived_latency_is_not_a_bottleneck_candidate() {
        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 repeated_stages_accumulate_and_milestones_do_not() {
        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)),
            "first audio describes startup; a later sentence must not lower it"
        );
    }

    #[test]
    fn summary_averages_across_turns() {
        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"));
    }
}