kcode-audio-ingress 0.7.1

//! Typed speaker-analysis validation and classifier orchestration.

use std::{collections::HashSet, sync::Arc};

use anyhow::{Context, ensure};
use chrono::{DateTime, Utc};
use kcode_speaker_extract::ExtractionOutcome;
use kcode_speaker_system::{Cohort, IdentifyEvidence, SpeechClassifier};
pub use kcode_speaker_system::{FeatureRow, ObservationKey};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

/// Exact classifier provider cohort component.
pub const CLASSIFIER_PROVIDER: &str = "google";
/// Exact classifier model cohort component.
pub const CLASSIFIER_MODEL: &str = "gemini-3.1-pro-preview";
/// Exact classifier prompt-version cohort component.
pub const CLASSIFIER_PROMPT_VERSION: &str = "gemini-transcript-speaker-24-freeform/2";
/// Exact classifier feature-schema cohort component.
pub const CLASSIFIER_SCHEMA_VERSION: &str = "gemini-speaker-24-normalized/1";

const READ_ONLY_IDENTIFY_THRESHOLD: f64 = 1e308;

/// One typed speaker row normalized from a raw Gemini response.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ParsedSpeaker {
    /// Exact chunk-local label, such as `Speaker 1`.
    pub local_label: String,
    /// Primary language identifier, absent when Gemini withheld a complete profile.
    pub primary_language: Option<String>,
    /// Validated 24-value row, absent when Gemini withheld a complete profile.
    pub feature_row: Option<FeatureRow>,
}

/// Complete normalized speaker structure for one raw Gemini result.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct ParsedChunk {
    /// Whether every substantive speaker supplied a complete feature profile.
    pub clip_valid: bool,
    /// Brief incompleteness reason, present exactly when `clip_valid` is false.
    pub clip_validity_reason: Option<String>,
    /// One row for every substantive chunk-local speaker.
    pub speakers: Vec<ParsedSpeaker>,
}

/// Classifier evidence retained for one chunk-local speaker.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct CandidateMapping {
    /// Best candidate's caller-owned full name.
    pub full_name: String,
    /// Best-candidate log-likelihood score relative to background at zero.
    pub score: f64,
    /// Runner-up log-likelihood score relative to background at zero.
    pub runner_up_score: Option<f64>,
}

/// Human resolution for one chunk-local speaker.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum SpeakerResolution {
    /// A known or newly entered speaker name that may be trained.
    Known {
        /// Exact human-approved full name.
        full_name: String,
    },
    /// A deliberately unidentified speaker that must not be trained.
    Unknown,
}

/// One deterministic classifier observation in a correction packet.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct CorrectionObservation {
    /// Chunk-local speaker label.
    pub local_label: String,
    /// Stable zero-based ordinal.
    pub speaker_ordinal: u32,
    /// Deterministic persisted training and correction key.
    pub observation_key: ObservationKey,
    /// Best available read-only classifier evidence.
    pub candidate: Option<CandidateMapping>,
    /// Human-approved resolution, absent until this chunk is signed off.
    pub resolution: Option<SpeakerResolution>,
}

/// One complete chunk in a recording-level correction packet.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct CorrectionChunk {
    /// Zero-based chronological chunk index.
    pub chunk_index: usize,
    /// Total recording chunk count.
    pub chunk_count: usize,
    /// Source-audio start in milliseconds.
    pub audio_start_ms: u64,
    /// Source-audio end in milliseconds.
    pub audio_end_ms: u64,
    /// Complete raw Gemini response without normalization.
    pub raw_gemini_response: String,
    /// Normalized feature structure from the single recording-wide GPT pass.
    pub parsed: ParsedChunk,
    /// Read-only classifier mappings for every substantive speaker.
    pub observations: Vec<CorrectionObservation>,
    /// Whether a human approved every resolution in this chunk.
    pub signed_off: bool,
}

/// Durable identity-confirmation lifecycle for a correction packet.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ConfirmationState {
    /// One or more chunks still await human signoff.
    Unconfirmed,
    /// Legacy state retained only so old packets remain decodable.
    AutomaticallyTrained,
    /// Every chunk has explicit human signoff.
    Confirmed,
}

/// Complete transport-neutral correction packet for one recording.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct CorrectionPacket {
    /// Stable recording UUID.
    pub recording_id: Uuid,
    /// Stable application user identifier associated with provider usage.
    pub user_id: String,
    /// Lowercase SHA-256 identity of the retained original bytes.
    pub sha256: String,
    /// Sanitized original filename.
    pub original_filename: String,
    /// Original retained file size in bytes.
    pub size_bytes: u64,
    /// Instant at which the recording began.
    pub recorded_at: DateTime<Utc>,
    /// Total chronological chunk count.
    pub chunk_count: usize,
    /// Every raw response, feature row, mapping, key, resolution, and interval.
    pub chunks: Vec<CorrectionChunk>,
    /// Current durable confirmation state.
    pub confirmation_state: ConfirmationState,
}

/// One observation-level human resolution.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct ObservationConfirmation {
    /// Exact deterministic observation key from the correction packet.
    pub observation_key: ObservationKey,
    /// Caller-confirmed known or unknown resolution.
    pub resolution: SpeakerResolution,
}

/// Exact speaker resolutions and signoff for one review chunk.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct ChunkConfirmation {
    /// Recording receiving the chunk signoff.
    pub recording_id: Uuid,
    /// Exact zero-based chunk index receiving signoff.
    pub chunk_index: usize,
    /// One resolution for every chunk-local observation, with no extras.
    pub observations: Vec<ObservationConfirmation>,
}

#[derive(Clone)]
pub(crate) struct ClassificationContext {
    pub(crate) recording_id: Uuid,
    pub(crate) user_id: String,
    pub(crate) sha256: String,
    pub(crate) original_filename: String,
    pub(crate) size_bytes: u64,
    pub(crate) recorded_at: DateTime<Utc>,
    pub(crate) classifier: Arc<SpeechClassifier>,
}

pub(crate) fn parsed_chunk_from_extraction(
    extraction: &ExtractionOutcome,
) -> anyhow::Result<ParsedChunk> {
    let (count, clip_valid, clip_validity_reason) = match extraction {
        ExtractionOutcome::Scored(scored) if scored.additional_speakers.is_empty() => {
            (scored.speakers.len(), true, None)
        }
        ExtractionOutcome::Scored(scored) => (
            scored.speakers.len() + scored.additional_speakers.len(),
            false,
            Some("One or more speakers lacked a complete feature profile.".to_owned()),
        ),
        ExtractionOutcome::Unscorable {
            reason,
            additional_speakers,
        } => (additional_speakers.len(), false, Some(reason.clone())),
    };
    let speakers = (0..count)
        .map(|ordinal| {
            let profile = match extraction {
                ExtractionOutcome::Scored(scored) => scored
                    .speakers
                    .iter()
                    .find(|profile| usize::from(profile.speaker_ordinal) == ordinal),
                ExtractionOutcome::Unscorable { .. } => None,
            };
            ParsedSpeaker {
                local_label: format!("Speaker {}", ordinal + 1),
                primary_language: profile.map(|value| value.primary_language.as_ref().to_owned()),
                feature_row: profile.map(|value| value.features),
            }
        })
        .collect();
    Ok(ParsedChunk {
        clip_valid,
        clip_validity_reason,
        speakers,
    })
}

pub(crate) fn classify_speakers(
    context: &ClassificationContext,
    chunk_index: usize,
    parsed: &ParsedChunk,
) -> anyhow::Result<Vec<CorrectionObservation>> {
    let mut observations = Vec::with_capacity(parsed.speakers.len());
    for (ordinal, speaker) in parsed.speakers.iter().enumerate() {
        let speaker_ordinal = u32::try_from(ordinal)
            .context("chunk has more speakers than the key schema supports")?;
        let candidate = match (
            speaker.primary_language.as_deref(),
            speaker.feature_row.as_ref(),
        ) {
            (Some(primary_language), Some(feature_row)) => {
                let probe_key = ObservationKey {
                    object_id: probe_object_id(context.recording_id, chunk_index),
                    piece_index: speaker_ordinal,
                };
                let outcome = context
                    .classifier
                    .identify(
                        probe_key.clone(),
                        cohort(primary_language),
                        *feature_row,
                        READ_ONLY_IDENTIFY_THRESHOLD,
                    )
                    .with_context(|| {
                        format!(
                            "read-only identity scoring failed for chunk {chunk_index} speaker {}",
                            speaker.local_label
                        )
                    })?;
                if outcome.speaker_id.is_some() {
                    context
                        .classifier
                        .delete(probe_key)
                        .context("removing an unexpectedly accepted read-only probe")?;
                }
                outcome.evidence.as_ref().map(candidate_mapping)
            }
            (None, None) => None,
            _ => anyhow::bail!("speaker profile is only partially present"),
        };
        observations.push(CorrectionObservation {
            local_label: speaker.local_label.clone(),
            speaker_ordinal,
            observation_key: ObservationKey {
                object_id: training_object_id(context.recording_id, chunk_index),
                piece_index: speaker_ordinal,
            },
            candidate,
            resolution: None,
        });
    }
    Ok(observations)
}

pub(crate) fn unclassified_observations(
    recording_id: Uuid,
    chunk_index: usize,
    parsed: &ParsedChunk,
) -> anyhow::Result<Vec<CorrectionObservation>> {
    parsed
        .speakers
        .iter()
        .enumerate()
        .map(|(ordinal, speaker)| {
            let speaker_ordinal = u32::try_from(ordinal)
                .context("chunk has more speakers than the key schema supports")?;
            Ok(CorrectionObservation {
                local_label: speaker.local_label.clone(),
                speaker_ordinal,
                observation_key: ObservationKey {
                    object_id: training_object_id(recording_id, chunk_index),
                    piece_index: speaker_ordinal,
                },
                candidate: None,
                resolution: None,
            })
        })
        .collect()
}

pub(crate) fn build_packet(
    context: &ClassificationContext,
    chunks: Vec<CorrectionChunk>,
) -> anyhow::Result<CorrectionPacket> {
    ensure!(!chunks.is_empty(), "correction packet has no chunks");
    let chunk_count = chunks.len();
    ensure!(
        chunks.iter().enumerate().all(|(index, chunk)| {
            chunk.chunk_index == index
                && chunk.chunk_count == chunk_count
                && chunk.audio_end_ms > chunk.audio_start_ms
                && chunk.observations.len() == chunk.parsed.speakers.len()
                && !chunk.signed_off
        }),
        "correction packet chunks are not one complete chronological plan"
    );
    Ok(CorrectionPacket {
        recording_id: context.recording_id,
        user_id: context.user_id.clone(),
        sha256: context.sha256.clone(),
        original_filename: context.original_filename.clone(),
        size_bytes: context.size_bytes,
        recorded_at: context.recorded_at,
        chunk_count,
        chunks,
        confirmation_state: ConfirmationState::Unconfirmed,
    })
}

pub(crate) fn validate_confirmation_coverage(
    packet: &CorrectionPacket,
    confirmation: &ChunkConfirmation,
) -> Result<(), String> {
    if confirmation.recording_id != packet.recording_id {
        return Err("Confirmation recording ID does not match the packet.".into());
    }
    let chunk = packet
        .chunks
        .get(confirmation.chunk_index)
        .filter(|chunk| chunk.chunk_index == confirmation.chunk_index)
        .ok_or_else(|| "Confirmation chunk does not exist.".to_owned())?;
    let known = chunk
        .observations
        .iter()
        .map(|observation| key_tuple(&observation.observation_key))
        .collect::<HashSet<_>>();
    let mut supplied = HashSet::new();
    for observation in &confirmation.observations {
        if let SpeakerResolution::Known { full_name } = &observation.resolution
            && (full_name.trim().is_empty() || full_name.chars().count() > 512)
        {
            return Err("Known speaker names must contain between 1 and 512 characters.".into());
        }
        if !supplied.insert(key_tuple(&observation.observation_key)) {
            return Err("Confirmation contains a duplicate observation key.".into());
        }
    }
    if supplied != known {
        return Err(
            "Chunk signoff must resolve every speaker exactly once, with no extras.".into(),
        );
    }
    Ok(())
}

pub(crate) fn apply_confirmations(
    classifier: &SpeechClassifier,
    packet: &mut CorrectionPacket,
    confirmation: &ChunkConfirmation,
    legacy_observation_keys: &HashSet<(String, u32)>,
) -> anyhow::Result<()> {
    validate_confirmation_coverage(packet, confirmation).map_err(anyhow::Error::msg)?;
    let chunk_index = confirmation.chunk_index;
    if packet.chunks[chunk_index].signed_off {
        let exact_retry = confirmation.observations.iter().all(|entry| {
            packet.chunks[chunk_index]
                .observations
                .iter()
                .find(|observation| observation.observation_key == entry.observation_key)
                .and_then(|observation| observation.resolution.as_ref())
                == Some(&entry.resolution)
        });
        ensure!(exact_retry, "signed-off chunk cannot be changed");
        return Ok(());
    }

    let assignments = confirmation
        .observations
        .iter()
        .map(|entry| {
            (
                key_tuple(&entry.observation_key),
                normalized(&entry.resolution),
            )
        })
        .collect::<std::collections::HashMap<_, _>>();
    let mut applied: Vec<ObservationKey> = Vec::new();
    for position in 0..packet.chunks[chunk_index].observations.len() {
        let observation = &packet.chunks[chunk_index].observations[position];
        let resolution = assignments
            .get(&key_tuple(&observation.observation_key))
            .context("validated confirmation assignment disappeared")?
            .clone();
        if !legacy_observation_keys.contains(&key_tuple(&observation.observation_key)) {
            let speaker = &packet.chunks[chunk_index].parsed.speakers[position];
            let result = match (
                &resolution,
                speaker.primary_language.as_deref(),
                speaker.feature_row,
            ) {
                (SpeakerResolution::Known { full_name }, Some(language), Some(row)) => classifier
                    .train(
                        observation.observation_key.clone(),
                        cohort(language),
                        row,
                        full_name.clone(),
                    )
                    .map(|_| ()),
                _ => classifier
                    .delete(observation.observation_key.clone())
                    .map(|_| ()),
            };
            if let Err(error) = result {
                let rollback = applied
                    .iter()
                    .rev()
                    .filter_map(|key| classifier.delete(key.clone()).err().map(|e| e.to_string()))
                    .collect::<Vec<_>>();
                if rollback.is_empty() {
                    anyhow::bail!("applying identity confirmations failed: {error}");
                }
                anyhow::bail!(
                    "applying identity confirmations failed: {error}; rollback also failed: {}",
                    rollback.join("; ")
                );
            }
            applied.push(observation.observation_key.clone());
        }
        packet.chunks[chunk_index].observations[position].resolution = Some(resolution);
    }
    packet.chunks[chunk_index].signed_off = true;
    if packet.chunks.iter().all(|chunk| chunk.signed_off) {
        packet.confirmation_state = ConfirmationState::Confirmed;
    }
    Ok(())
}

pub(crate) fn restore_packet_training(
    classifier: &SpeechClassifier,
    packet: &CorrectionPacket,
    legacy_observation_keys: &HashSet<(String, u32)>,
) -> Vec<String> {
    let mut errors = Vec::new();
    for chunk in packet.chunks.iter().rev() {
        for observation in chunk.observations.iter().rev() {
            if legacy_observation_keys.contains(&key_tuple(&observation.observation_key)) {
                continue;
            }
            let speaker = chunk
                .parsed
                .speakers
                .get(observation.speaker_ordinal as usize);
            let result = match (observation.resolution.as_ref(), speaker) {
                (
                    Some(SpeakerResolution::Known { full_name }),
                    Some(ParsedSpeaker {
                        primary_language: Some(language),
                        feature_row: Some(row),
                        ..
                    }),
                ) => classifier
                    .train(
                        observation.observation_key.clone(),
                        cohort(language),
                        *row,
                        full_name.clone(),
                    )
                    .map(|_| ()),
                _ => classifier
                    .delete(observation.observation_key.clone())
                    .map(|_| ()),
            };
            if let Err(error) = result {
                errors.push(error.to_string());
            }
        }
    }
    errors
}

pub(crate) fn training_object_id(recording_id: Uuid, chunk_index: usize) -> String {
    format!("kcode-audio-ingress/recording/{recording_id}/chunk/{chunk_index}")
}

fn probe_object_id(recording_id: Uuid, chunk_index: usize) -> String {
    format!("kcode-audio-ingress/probe/{recording_id}/chunk/{chunk_index}")
}

fn cohort(primary_language: &str) -> Cohort {
    Cohort {
        provider: CLASSIFIER_PROVIDER.into(),
        model: CLASSIFIER_MODEL.into(),
        prompt_version: CLASSIFIER_PROMPT_VERSION.into(),
        schema_version: CLASSIFIER_SCHEMA_VERSION.into(),
        primary_language: primary_language.into(),
    }
}

fn candidate_mapping(evidence: &IdentifyEvidence) -> CandidateMapping {
    CandidateMapping {
        full_name: evidence.best.speaker_id.clone(),
        score: -evidence.best.cost,
        runner_up_score: evidence.runner_up.as_ref().map(|candidate| -candidate.cost),
    }
}

fn normalized(resolution: &SpeakerResolution) -> SpeakerResolution {
    match resolution {
        SpeakerResolution::Known { full_name } => SpeakerResolution::Known {
            full_name: full_name.trim().to_owned(),
        },
        SpeakerResolution::Unknown => SpeakerResolution::Unknown,
    }
}

fn key_tuple(key: &ObservationKey) -> (String, u32) {
    (key.object_id.clone(), key.piece_index)
}

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

    fn chunk(recording_id: Uuid, index: usize) -> CorrectionChunk {
        let key = ObservationKey {
            object_id: training_object_id(recording_id, index),
            piece_index: 0,
        };
        CorrectionChunk {
            chunk_index: index,
            chunk_count: 2,
            audio_start_ms: index as u64 * 1_000,
            audio_end_ms: (index as u64 + 1) * 1_000,
            raw_gemini_response: "Speaker 1: hello".into(),
            parsed: ParsedChunk {
                clip_valid: true,
                clip_validity_reason: None,
                speakers: vec![ParsedSpeaker {
                    local_label: "Speaker 1".into(),
                    primary_language: Some("eng".into()),
                    feature_row: Some(FeatureRow::new([50; 24]).unwrap()),
                }],
            },
            observations: vec![CorrectionObservation {
                local_label: "Speaker 1".into(),
                speaker_ordinal: 0,
                observation_key: key,
                candidate: None,
                resolution: None,
            }],
            signed_off: false,
        }
    }

    #[test]
    fn per_chunk_signoff_trains_known_and_skips_unknown_speakers() {
        let recording_id = Uuid::new_v4();
        let path = std::env::temp_dir().join(format!("audio-ingress-{recording_id}.sqlite3"));
        let classifier = SpeechClassifier::open(&path).unwrap();
        let mut packet = CorrectionPacket {
            recording_id,
            user_id: "user".into(),
            sha256: "a".repeat(64),
            original_filename: "voice.wav".into(),
            size_bytes: 1,
            recorded_at: Utc::now(),
            chunk_count: 2,
            chunks: vec![chunk(recording_id, 0), chunk(recording_id, 1)],
            confirmation_state: ConfirmationState::Unconfirmed,
        };
        let keys = packet
            .chunks
            .iter()
            .map(|chunk| chunk.observations[0].observation_key.clone())
            .collect::<Vec<_>>();
        apply_confirmations(
            &classifier,
            &mut packet,
            &ChunkConfirmation {
                recording_id,
                chunk_index: 0,
                observations: vec![ObservationConfirmation {
                    observation_key: keys[0].clone(),
                    resolution: SpeakerResolution::Known {
                        full_name: "David Example".into(),
                    },
                }],
            },
            &HashSet::new(),
        )
        .unwrap();
        assert_eq!(packet.confirmation_state, ConfirmationState::Unconfirmed);
        apply_confirmations(
            &classifier,
            &mut packet,
            &ChunkConfirmation {
                recording_id,
                chunk_index: 1,
                observations: vec![ObservationConfirmation {
                    observation_key: keys[1].clone(),
                    resolution: SpeakerResolution::Unknown,
                }],
            },
            &HashSet::new(),
        )
        .unwrap();
        assert_eq!(packet.confirmation_state, ConfirmationState::Confirmed);
        assert_eq!(classifier.known_speakers().unwrap(), vec!["David Example"]);
        drop(classifier);
        let _ = std::fs::remove_file(path);
    }
}