kcode-k1-audio-classification-event-format 0.1.0

use serde::{Deserialize, Serialize};

pub use kcode_k1_audio_classification_format_error::FormatError;
pub use kcode_k1_transaction_id::TxId;
pub use kcode_speaker_v3_analysis::{ExecutedAnalysis, FeatureVector24, LocalSpeakerLabel};

const EVENT_VERSION: u8 = 4;
const QUEUE_TAG: u8 = 1;
const TRANSCRIPTION_COMPLETE_TAG: u8 = 2;
const FAILED_TAG: u8 = 3;
const DISCARDED_TAG: u8 = 4;
const LABEL_CONFIRMATION_TAG: u8 = 5;
const PROGRESS_TAG: u8 = 6;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FragmentStageV1 {
    Queue,
    Transcript,
    SpeakerLabels,
    SpeakerFeatures,
    Structuring,
    LabelConfirmation,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct QueueV2 {
    #[serde(with = "txid_serde")]
    pub audio_object_id: TxId,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ProgressV1 {
    #[serde(with = "txid_serde")]
    pub fragment_id: TxId,
    pub update: ProgressUpdateV1,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProgressUpdateV1 {
    LlmJobStarted {
        sequence: u64,
        stage: FragmentStageV1,
        name: String,
    },
    LlmJobSucceeded {
        sequence: u64,
    },
    LlmJobFailed {
        sequence: u64,
        error: String,
    },
    StageCompleted {
        stage: FragmentStageV1,
    },
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TranscriptionCompleteV1 {
    #[serde(with = "txid_serde")]
    pub fragment_id: TxId,
    pub analysis: ExecutedAnalysis,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FailedV2 {
    #[serde(with = "txid_serde")]
    pub fragment_id: TxId,
    pub stage: FragmentStageV1,
    pub llm_job_sequence: Option<u64>,
    pub error: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DiscardedV2 {
    #[serde(with = "txid_serde")]
    pub fragment_id: TxId,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SpeakerLabelV1 {
    pub speaker: LocalSpeakerLabel,
    pub person_id: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LabelConfirmationV1 {
    #[serde(with = "txid_serde")]
    pub fragment_id: TxId,
    #[serde(with = "txid_serde")]
    pub interim_txid: TxId,
    pub speakers: Vec<SpeakerLabelV1>,
}

#[allow(clippy::large_enum_variant)]
#[derive(Debug, Clone, PartialEq)]
pub enum AudioClassificationEventV3 {
    Queue(QueueV2),
    Progress(ProgressV1),
    TranscriptionComplete(TranscriptionCompleteV1),
    Failed(FailedV2),
    Discarded(DiscardedV2),
    LabelConfirmation(LabelConfirmationV1),
}

pub fn encode_event(event: &AudioClassificationEventV3) -> Result<Vec<u8>, FormatError> {
    let (tag, body) = match event {
        AudioClassificationEventV3::Queue(value) => (QUEUE_TAG, postcard::to_allocvec(value)),
        AudioClassificationEventV3::TranscriptionComplete(value) => {
            (TRANSCRIPTION_COMPLETE_TAG, postcard::to_allocvec(value))
        }
        AudioClassificationEventV3::Failed(value) => (FAILED_TAG, postcard::to_allocvec(value)),
        AudioClassificationEventV3::Discarded(value) => {
            (DISCARDED_TAG, postcard::to_allocvec(value))
        }
        AudioClassificationEventV3::LabelConfirmation(value) => {
            (LABEL_CONFIRMATION_TAG, postcard::to_allocvec(value))
        }
        AudioClassificationEventV3::Progress(value) => (PROGRESS_TAG, postcard::to_allocvec(value)),
    };
    let body = body.map_err(|_| FormatError::InvalidEventBody)?;
    let mut output = Vec::with_capacity(2 + body.len());
    output.extend_from_slice(&[EVENT_VERSION, tag]);
    output.extend_from_slice(&body);
    Ok(output)
}

pub fn decode_event(bytes: &[u8]) -> Result<AudioClassificationEventV3, FormatError> {
    if bytes.len() < 2 {
        return Err(FormatError::Truncated);
    }
    if bytes[0] != EVENT_VERSION {
        return Err(FormatError::UnsupportedVersion(bytes[0]));
    }
    match bytes[1] {
        QUEUE_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Queue),
        TRANSCRIPTION_COMPLETE_TAG => {
            decode_body(&bytes[2..]).map(AudioClassificationEventV3::TranscriptionComplete)
        }
        FAILED_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Failed),
        DISCARDED_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Discarded),
        LABEL_CONFIRMATION_TAG => {
            decode_body(&bytes[2..]).map(AudioClassificationEventV3::LabelConfirmation)
        }
        PROGRESS_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Progress),
        tag => Err(FormatError::UnknownEventTag(tag)),
    }
}

fn decode_body<T: for<'de> Deserialize<'de>>(bytes: &[u8]) -> Result<T, FormatError> {
    let (value, remaining) =
        postcard::take_from_bytes(bytes).map_err(|_| FormatError::InvalidEventBody)?;
    if !remaining.is_empty() {
        return Err(FormatError::TrailingBytes);
    }
    Ok(value)
}

mod txid_serde {
    use super::TxId;
    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    pub fn serialize<S: Serializer>(value: &TxId, serializer: S) -> Result<S::Ok, S::Error> {
        value.into_bytes().serialize(serializer)
    }

    pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<TxId, D::Error> {
        Ok(TxId::from_bytes(<[u8; 12]>::deserialize(deserializer)?))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use kcode_speaker_v3_analysis::{
        AnalysisEnvelope, GeminiCohort, OggAudioMetadata, StructuredAnalysis, StructuredSpeaker,
        StructurerProvenance,
    };
    use serde::Serialize;

    fn txid(seed: u8) -> TxId {
        let mut bytes = [0; 12];
        for (index, value) in bytes.iter_mut().enumerate() {
            *value = seed.wrapping_add(index as u8);
        }
        TxId::from_bytes(bytes)
    }

    fn analysis() -> ExecutedAnalysis {
        let mut ogg = vec![0; 29];
        ogg[..4].copy_from_slice(b"OggS");
        ogg[26] = 1;
        ogg[27] = 1;
        ExecutedAnalysis {
            envelope: AnalysisEnvelope {
                audio: OggAudioMetadata::from_bytes(&ogg, 1_250, Some("sample.ogg".into()))
                    .expect("Ogg metadata"),
                analysis: StructuredAnalysis {
                    transcript: "Speaker 1".into(),
                    speakers: vec![StructuredSpeaker {
                        speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
                        language: "English".into(),
                        features: FeatureVector24 {
                            median_f0_hz: Some(182.5),
                            ..FeatureVector24::default()
                        },
                        features_usable_for_training: true,
                    }],
                },
                gemini: GeminiCohort::new("gemini-3.1-pro"),
                structurer: StructurerProvenance::new("terra-structurer"),
            },
            label_extractor: StructurerProvenance::new("terra-labeler"),
        }
    }

    fn assert_golden<T: Serialize>(event: AudioClassificationEventV3, tag: u8, body: &T) {
        let mut expected = vec![EVENT_VERSION, tag];
        expected.extend(postcard::to_allocvec(body).expect("serialize golden body"));
        assert_eq!(encode_event(&event), Ok(expected.clone()));
        assert_eq!(decode_event(&expected), Ok(event));
    }

    #[test]
    fn all_event_variants_match_v4_golden_envelopes() {
        let queue = QueueV2 {
            audio_object_id: txid(1),
        };
        assert_golden(
            AudioClassificationEventV3::Queue(queue.clone()),
            QUEUE_TAG,
            &queue,
        );

        let complete = TranscriptionCompleteV1 {
            fragment_id: txid(2),
            analysis: analysis(),
        };
        assert_golden(
            AudioClassificationEventV3::TranscriptionComplete(complete.clone()),
            TRANSCRIPTION_COMPLETE_TAG,
            &complete,
        );

        let failed = FailedV2 {
            fragment_id: txid(3),
            stage: FragmentStageV1::SpeakerFeatures,
            llm_job_sequence: Some(7),
            error: "failure".into(),
        };
        assert_golden(
            AudioClassificationEventV3::Failed(failed.clone()),
            FAILED_TAG,
            &failed,
        );

        let discarded = DiscardedV2 {
            fragment_id: txid(4),
        };
        assert_golden(
            AudioClassificationEventV3::Discarded(discarded.clone()),
            DISCARDED_TAG,
            &discarded,
        );

        let confirmation = LabelConfirmationV1 {
            fragment_id: txid(5),
            interim_txid: txid(6),
            speakers: vec![SpeakerLabelV1 {
                speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
                person_id: "person-1".into(),
            }],
        };
        assert_golden(
            AudioClassificationEventV3::LabelConfirmation(confirmation.clone()),
            LABEL_CONFIRMATION_TAG,
            &confirmation,
        );

        let progress = ProgressV1 {
            fragment_id: txid(7),
            update: ProgressUpdateV1::LlmJobStarted {
                sequence: 9,
                stage: FragmentStageV1::Transcript,
                name: "transcribe".into(),
            },
        };
        assert_golden(
            AudioClassificationEventV3::Progress(progress.clone()),
            PROGRESS_TAG,
            &progress,
        );
    }

    #[test]
    fn every_progress_update_roundtrips() {
        let updates = [
            ProgressUpdateV1::LlmJobStarted {
                sequence: 1,
                stage: FragmentStageV1::Transcript,
                name: "transcript".into(),
            },
            ProgressUpdateV1::LlmJobSucceeded { sequence: 1 },
            ProgressUpdateV1::LlmJobFailed {
                sequence: 2,
                error: "failed".into(),
            },
            ProgressUpdateV1::StageCompleted {
                stage: FragmentStageV1::Structuring,
            },
        ];
        for update in updates {
            let event = AudioClassificationEventV3::Progress(ProgressV1 {
                fragment_id: txid(8),
                update,
            });
            let bytes = encode_event(&event).expect("encode progress");
            assert_eq!(decode_event(&bytes), Ok(event));
        }
    }

    #[test]
    fn canonical_txid_encoding_is_twelve_raw_bytes() {
        let id = txid(0);
        let event = AudioClassificationEventV3::Queue(QueueV2 {
            audio_object_id: id,
        });
        let mut expected = vec![EVENT_VERSION, QUEUE_TAG];
        expected.extend_from_slice(id.as_bytes());
        assert_eq!(encode_event(&event), Ok(expected));
    }

    #[test]
    fn current_speaker_analysis_serialization_is_embedded_unchanged() {
        let fragment_id = txid(12);
        let analysis = analysis();
        let expected_analysis = postcard::to_allocvec(&analysis).expect("serialize analysis");
        let bytes = encode_event(&AudioClassificationEventV3::TranscriptionComplete(
            TranscriptionCompleteV1 {
                fragment_id,
                analysis,
            },
        ))
        .expect("encode complete event");
        assert_eq!(&bytes[..2], &[EVENT_VERSION, TRANSCRIPTION_COMPLETE_TAG]);
        assert_eq!(&bytes[2..14], fragment_id.as_bytes());
        assert_eq!(&bytes[14..], expected_analysis);
    }

    #[test]
    fn malformed_unknown_version_and_trailing_inputs_are_rejected() {
        assert_eq!(decode_event(&[]), Err(FormatError::Truncated));
        assert_eq!(decode_event(&[EVENT_VERSION]), Err(FormatError::Truncated));
        assert_eq!(
            decode_event(&[3, QUEUE_TAG]),
            Err(FormatError::UnsupportedVersion(3))
        );
        assert_eq!(
            decode_event(&[EVENT_VERSION, 7]),
            Err(FormatError::UnknownEventTag(7))
        );
        assert_eq!(
            decode_event(&[EVENT_VERSION, QUEUE_TAG]),
            Err(FormatError::InvalidEventBody)
        );

        let event = AudioClassificationEventV3::Discarded(DiscardedV2 {
            fragment_id: txid(9),
        });
        let mut bytes = encode_event(&event).expect("encode discarded");
        bytes.push(0);
        assert_eq!(decode_event(&bytes), Err(FormatError::TrailingBytes));
    }
}