//! Разбор по голосам после распознавания.
//!
//! ```text
//! реплики Whisper ─► отпечаток голоса каждой (кусок звука с диска) ─► группы ─► «Спикер 1», «Спикер 2»…
//! ```
//!
//! Разбираются дорожки, где может говорить больше одного человека: звук
//! компьютера (собеседники) и файл целиком. Микрофон — всегда «Я».
//!
//! Реплики с отпечатками лежат рядом с записью (`voices.json`): так разбиение
//! пересобирается на другое число говорящих за секунду, без новой расшифровки.
//! Отпечаток — не голос: из него звук не восстановить, он нужен только, чтобы
//! сравнить две реплики этой записи между собой.

use std::io;
use std::sync::Arc;

use base64::Engine;
use meet_audio::pcm::ms_to_samples;
use meet_voices::{Count, Extractor};
use serde::{Deserialize, Serialize};
use tokio_util::sync::CancellationToken;

use crate::transcribe::{Samples, UNRECOGNIZED};
use crate::transcript::{Segment, Speaker};

/// Короче — отпечаток не снимается: модели не хватает звука.
const MIN_EMBED_MS: u64 = 500;
/// Длиннее — берётся середина: хватает, чтобы узнать голос, и быстрее.
const MAX_EMBED_MS: u64 = 12_000;

/// Реплики записи с отпечатками — то, из чего собирается разбиение.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Voices {
    pub segments: Vec<Stored>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Stored {
    pub start_ms: u64,
    pub end_ms: u64,
    /// Дорожка, с которой пришла реплика.
    pub track: Speaker,
    pub text: String,
    /// Отпечаток (f32, little-endian) в base64; пусто — не снят.
    #[serde(default)]
    pub voice: String,
}

/// Эту дорожку надо разбирать по голосам.
pub fn diarized(track: &Speaker) -> bool {
    matches!(track, Speaker::Others | Speaker::Unknown)
}

fn encode(v: &[f32]) -> String {
    let bytes: Vec<u8> = v.iter().flat_map(|x| x.to_le_bytes()).collect();
    base64::engine::general_purpose::STANDARD.encode(bytes)
}

fn decode(s: &str) -> Vec<f32> {
    base64::engine::general_purpose::STANDARD
        .decode(s)
        .map(|b| {
            b.as_chunks::<4>()
                .0
                .iter()
                .map(|c| f32::from_le_bytes(*c))
                .collect()
        })
        .unwrap_or_default()
}

/// Разбор отменил человек.
#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("Расшифровка отменена")]
    Cancelled,
    #[error("Не читается звук: {0}")]
    Read(#[from] io::Error),
}

/// Снимает отпечатки реплик разбираемых дорожек. Работает долго (минуты на
/// двухчасовую встречу) — зовётся из отдельного потока; `progress(готово,
/// всего)` — после каждой реплики.
pub fn embed(
    segments: &[Segment],
    sources: &[(Speaker, Arc<dyn Samples>)],
    extractor: &mut Extractor,
    cancel: &CancellationToken,
    progress: impl Fn(usize, usize),
) -> Result<Voices, Error> {
    let todo = segments
        .iter()
        .filter(|s| diarized(&s.speaker) && s.text != UNRECOGNIZED)
        .count();
    let mut done = 0;
    progress(0, todo);
    let mut out = Vec::with_capacity(segments.len());
    for s in segments {
        let mut voice = String::new();
        if diarized(&s.speaker) && s.text != UNRECOGNIZED {
            if cancel.is_cancelled() {
                return Err(Error::Cancelled);
            }
            let len = s.end_ms.saturating_sub(s.start_ms);
            if len >= MIN_EMBED_MS {
                if let Some((_, source)) = sources.iter().find(|(t, _)| t == &s.speaker) {
                    let (from, len) = if len > MAX_EMBED_MS {
                        (s.start_ms + (len - MAX_EMBED_MS) / 2, MAX_EMBED_MS)
                    } else {
                        (s.start_ms, len)
                    };
                    let samples = source.read(ms_to_samples(from), ms_to_samples(len))?;
                    let v = extractor.embed(&samples);
                    if !v.is_empty() {
                        voice = encode(&v);
                    }
                }
            }
            done += 1;
            progress(done, todo);
        }
        out.push(Stored {
            start_ms: s.start_ms,
            end_ms: s.end_ms,
            track: s.speaker.clone(),
            text: s.text.clone(),
            voice,
        });
    }
    Ok(Voices { segments: out })
}

/// Реплики с говорящими: разбираемые дорожки — «Спикер N» (N с единицы, по
/// порядку первого слова), микрофон — «Я».
pub fn label(voices: &Voices, count: Count) -> Vec<Segment> {
    let picked: Vec<usize> = (0..voices.segments.len())
        .filter(|&i| diarized(&voices.segments[i].track))
        .collect();
    let embeddings: Vec<Vec<f32>> = picked
        .iter()
        .map(|&i| decode(&voices.segments[i].voice))
        .collect();
    let seconds: Vec<f32> = picked
        .iter()
        .map(|&i| {
            let s = &voices.segments[i];
            s.end_ms.saturating_sub(s.start_ms) as f32 / 1000.0
        })
        .collect();
    let labels = meet_voices::assign(&embeddings, &seconds, count);
    let mut speaker: Vec<Speaker> = voices.segments.iter().map(|s| s.track.clone()).collect();
    for (&i, n) in picked.iter().zip(labels) {
        speaker[i] = Speaker::Voice { n: n + 1 };
    }
    voices
        .segments
        .iter()
        .zip(speaker)
        .map(|(s, speaker)| Segment {
            start_ms: s.start_ms,
            end_ms: s.end_ms,
            speaker,
            text: s.text.clone(),
        })
        .collect()
}

/// Сколько разных «Спикеров» в разбиении.
pub fn voice_count(segments: &[Segment]) -> usize {
    segments
        .iter()
        .filter_map(|s| match s.speaker {
            Speaker::Voice { n } => Some(n),
            _ => None,
        })
        .max()
        .unwrap_or(0)
}

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

    fn stored(ms: u64, track: Speaker, voice: &[f32]) -> Stored {
        Stored {
            start_ms: ms,
            end_ms: ms + 3_000,
            track,
            text: format!("реплика {ms}"),
            voice: if voice.is_empty() {
                String::new()
            } else {
                encode(voice)
            },
        }
    }

    #[test]
    fn voice_survives_the_round_trip() {
        let v = [0.25f32, -1.5, 3.0e-7];
        assert_eq!(decode(&encode(&v)), v);
        assert!(decode("не base64").is_empty());
    }

    #[test]
    fn my_track_stays_me_and_others_get_numbers() {
        // Два голоса собеседников — два направления с общим сдвигом записи.
        let a = [1.0, 0.1, 0.0, 0.5];
        let b = [0.1, 1.0, 0.0, 0.5];
        let voices = Voices {
            segments: vec![
                stored(0, Speaker::Me, &[]),
                stored(3_000, Speaker::Others, &b),
                stored(6_000, Speaker::Others, &a),
                stored(9_000, Speaker::Others, &b),
                stored(12_000, Speaker::Others, &[]),
                stored(15_000, Speaker::Others, &a),
            ],
        };
        let got: Vec<Speaker> = label(&voices, Count::Fixed(2))
            .into_iter()
            .map(|s| s.speaker)
            .collect();
        let v = |n| Speaker::Voice { n };
        assert_eq!(got, [Speaker::Me, v(1), v(2), v(1), v(1), v(2)]);
        let one = label(&voices, Count::Fixed(1));
        assert_eq!(voice_count(&one), 1);
    }

    #[test]
    fn stored_voices_read_old_files_without_voice() {
        let json = r#"{"segments":[{"start_ms":0,"end_ms":10,"track":{"kind":"unknown"},"text":"Да."}]}"#;
        let v: Voices = serde_json::from_str(json).unwrap();
        assert_eq!(v.segments[0].voice, "");
        assert_eq!(
            label(&v, Count::Auto { threshold: 0.0 })[0].speaker,
            Speaker::Voice { n: 1 }
        );
    }

    /// Весь путь на настоящей записи, без OpenRouter: звук → реплики по паузам
    /// (вместо Whisper) → отпечатки → разбиение → заметка. По запросу:
    /// `MEET_VOICES_AUDIO=<файл> MEET_VOICES_MODEL=models/voices.onnx cargo test
    /// --release -p skillok-meet real_recording -- --ignored --nocapture`.
    #[test]
    #[ignore]
    fn real_recording() {
        use meet_audio::chunker::{Chunker, Config};
        use meet_audio::pcm::{samples_to_ms, WavFile, WavWriter};
        let audio = std::env::var("MEET_VOICES_AUDIO").expect("MEET_VOICES_AUDIO");
        let model = std::env::var("MEET_VOICES_MODEL").unwrap_or_else(|_| "../models/voices.onnx".into());
        let dir = tempfile::tempdir().unwrap();
        let wav = dir.path().join("track.wav");
        {
            let pcm = meet_audio::decode::decode_file(std::path::Path::new(&audio)).unwrap();
            let mut w = WavWriter::create(&wav).unwrap();
            w.write(&pcm).unwrap();
        }
        let file = WavFile::open(&wav).unwrap();
        // Реплики, как их примерно режет Whisper: между паузами, 1–10 с.
        let mut c = Chunker::new(Config {
            min_ms: 1_000,
            max_ms: 10_000,
            pause_ms: 300,
            min_speech_ms: 300,
            ..Config::default()
        });
        let mut segments = Vec::new();
        let block = 160_000;
        let mut pos = 0;
        while pos < file.len() {
            let part = file.read(pos, block).unwrap();
            pos += part.len();
            for ch in c.push(samples_to_ms(pos), &part) {
                segments.push(Segment {
                    start_ms: ch.start_ms,
                    end_ms: ch.start_ms + ch.duration_ms(),
                    speaker: Speaker::Unknown,
                    text: format!("реплика {}", segments.len()),
                });
            }
        }
        let sources: Vec<(Speaker, Arc<dyn Samples>)> = vec![(Speaker::Unknown, Arc::new(file))];
        let threads = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(4);
        let mut ex = Extractor::open(std::path::Path::new(&model), threads).unwrap();
        let started = std::time::Instant::now();
        let voices = embed(&segments, &sources, &mut ex, &CancellationToken::new(), |_, _| {}).unwrap();
        let took = started.elapsed().as_secs_f32();
        let auto = label(
            &voices,
            Count::Auto {
                threshold: meet_voices::DEFAULT_THRESHOLD,
            },
        );
        let mut talk = std::collections::BTreeMap::new();
        for s in &auto {
            *talk.entry(s.speaker.label().to_string()).or_insert(0u64) += s.end_ms - s.start_ms;
        }
        let total: u64 = talk.values().sum();
        println!("{} реплик, отпечатки за {took:.1} с", segments.len());
        for (who, ms) in &talk {
            println!("  {who}: {} % речи", ms * 100 / total.max(1));
        }
        let md = crate::transcript::to_markdown(&auto);
        println!("{}", md.lines().take(24).collect::<Vec<_>>().join("\n"));
        assert!(voice_count(&auto) >= 1);
        assert_eq!(voice_count(&label(&voices, Count::Fixed(2))), 2);
    }
}