kcode-k1-audio-classification-format 0.6.1

use kcode_k1_audio_classification_format as facade;
use kcode_speaker_v3_analysis as current;
use kcode_speaker_v3_analysis_0_2 as previous;
use serde::Serialize;

#[derive(Serialize)]
struct PreviousTranscriptionCompleteV1 {
    fragment_id: [u8; 12],
    analysis: previous::ExecutedAnalysis,
}

fn txid(seed: u8) -> facade::TxId {
    facade::TxId::from_bytes(std::array::from_fn(|index| seed.wrapping_add(index as u8)))
}

fn person(seed: u8) -> facade::PersonId {
    facade::PersonId::from_tx_id(txid(seed))
}

macro_rules! analysis {
    ($owner:ident, $message:literal) => {{
        let mut ogg = vec![0; 29];
        ogg[..4].copy_from_slice(b"OggS");
        ogg[26] = 1;
        ogg[27] = 1;
        $owner::ExecutedAnalysis {
            envelope: $owner::AnalysisEnvelope {
                audio: $owner::OggAudioMetadata::from_bytes(&ogg, 1_250, Some("sample.ogg".into()))
                    .expect($message),
                analysis: $owner::StructuredAnalysis {
                    transcript: "Speaker 1: hello".into(),
                    speakers: vec![$owner::StructuredSpeaker {
                        speaker: $owner::LocalSpeakerLabel::new(1).expect("speaker label"),
                        language: "English".into(),
                        features: $owner::FeatureVector24::default(),
                        features_usable_for_training: true,
                    }],
                },
                gemini: $owner::GeminiCohort {
                    model_id: "gemini".into(),
                    transcript_prompt_revision: "transcript-r1".into(),
                    feature_prompt_revisions: [
                        "feature-r1".into(),
                        "feature-r2".into(),
                        "feature-r3".into(),
                    ],
                    feature_schema_revision: "schema-r1".into(),
                },
                structurer: $owner::StructurerProvenance {
                    model_id: "terra".into(),
                    prompt_revision: "structuring-r1".into(),
                },
            },
            label_extractor: $owner::StructurerProvenance {
                model_id: "terra".into(),
                prompt_revision: "labels-r1".into(),
            },
        }
    }};
}

fn current_analysis() -> current::ExecutedAnalysis {
    analysis!(current, "current Ogg metadata")
}

fn previous_analysis() -> previous::ExecutedAnalysis {
    analysis!(previous, "previous Ogg metadata")
}

fn frozen_event<T: Serialize>(tag: u8, value: &T) -> Vec<u8> {
    let mut bytes = vec![5, tag];
    bytes.extend(postcard::to_allocvec(value).expect("serialize frozen event body"));
    bytes
}

fn assert_event<T: Serialize>(tag: u8, event: facade::AudioClassificationEventV3, value: &T) {
    let bytes = frozen_event(tag, value);
    assert_eq!(facade::encode_event(&event), Ok(bytes.clone()));
    assert_eq!(facade::decode_event(&bytes), Ok(event));
}

fn append_length(bytes: &mut Vec<u8>, length: usize) {
    bytes.extend_from_slice(&(length as u64).to_le_bytes());
}

fn append_bytes(bytes: &mut Vec<u8>, value: &[u8]) {
    append_length(bytes, value.len());
    bytes.extend_from_slice(value);
}

fn append_string(bytes: &mut Vec<u8>, value: &str) {
    append_bytes(bytes, value.as_bytes());
}

fn previous_fragment_prefix(
    kind: u8,
    analysis: facade::TxId,
    confirmation: Option<facade::TxId>,
) -> Vec<u8> {
    let mut bytes = vec![0; 48];
    bytes[0] = 2;
    bytes[1] = kind;
    bytes[16..28].copy_from_slice(analysis.as_bytes());
    if let Some(confirmation) = confirmation {
        bytes[32..44].copy_from_slice(confirmation.as_bytes());
    }
    bytes
}

fn previous_staged_fragment_bytes(features: &previous::FeatureVector24) -> Vec<u8> {
    let mut bytes = previous_fragment_prefix(1, txid(20), None);
    append_string(&mut bytes, "hello");
    append_length(&mut bytes, 1);
    append_string(
        &mut bytes,
        &previous::LocalSpeakerLabel::new(1)
            .expect("previous speaker label")
            .to_string(),
    );
    append_string(&mut bytes, "English");
    append_bytes(
        &mut bytes,
        &postcard::to_allocvec(features).expect("serialize previous features"),
    );
    bytes.push(1);
    bytes
}

fn previous_final_fragment_bytes(
    features: &previous::FeatureVector24,
    known: facade::PersonId,
) -> Vec<u8> {
    let mut bytes = previous_fragment_prefix(2, txid(30), Some(txid(40)));
    append_string(&mut bytes, "hello");
    append_length(&mut bytes, 1);
    append_string(
        &mut bytes,
        &previous::LocalSpeakerLabel::new(1)
            .expect("previous speaker label")
            .to_string(),
    );
    bytes.push(1);
    bytes.extend_from_slice(known.as_tx_id().as_bytes());
    append_string(&mut bytes, "English");
    append_bytes(
        &mut bytes,
        &postcard::to_allocvec(features).expect("serialize previous features"),
    );
    bytes.push(0);
    bytes
}

#[test]
fn all_public_paths_compile_and_non_identity_event_bytes_remain_frozen() {
    let speaker = facade::LocalSpeakerLabel::new(1).expect("speaker label");
    let queue = facade::QueueV2 {
        audio_object_id: txid(1),
    };
    assert_event(
        1,
        facade::AudioClassificationEventV3::Queue(queue.clone()),
        &queue,
    );
    let updates = [
        facade::ProgressUpdateV1::LlmJobStarted {
            sequence: 1,
            stage: facade::FragmentStageV1::Transcript,
            name: "transcribe".into(),
        },
        facade::ProgressUpdateV1::LlmJobSucceeded { sequence: 2 },
        facade::ProgressUpdateV1::LlmJobFailed {
            sequence: 3,
            error: "failed".into(),
        },
        facade::ProgressUpdateV1::StageCompleted {
            stage: facade::FragmentStageV1::Structuring,
        },
    ];
    for update in updates {
        let progress = facade::ProgressV1 {
            fragment_id: txid(2),
            update,
        };
        assert_event(
            6,
            facade::AudioClassificationEventV3::Progress(progress.clone()),
            &progress,
        );
    }
    let transcription = facade::TranscriptionCompleteV1 {
        fragment_id: txid(3),
        analysis: current_analysis(),
    };
    assert_event(
        2,
        facade::AudioClassificationEventV3::TranscriptionComplete(transcription.clone()),
        &transcription,
    );
    let failed = facade::FailedV2 {
        fragment_id: txid(4),
        stage: facade::FragmentStageV1::SpeakerFeatures,
        llm_job_sequence: Some(7),
        error: "failure".into(),
    };
    assert_event(
        3,
        facade::AudioClassificationEventV3::Failed(failed.clone()),
        &failed,
    );
    let discarded = facade::DiscardedV2 {
        fragment_id: txid(5),
    };
    assert_event(
        4,
        facade::AudioClassificationEventV3::Discarded(discarded.clone()),
        &discarded,
    );
    let confirmation = facade::LabelConfirmationV1 {
        fragment_id: txid(6),
        interim_txid: txid(7),
        speakers: vec![facade::SpeakerLabelV1 {
            speaker,
            person_id: Some(person(70)),
        }],
    };
    assert_event(
        5,
        facade::AudioClassificationEventV3::LabelConfirmation(confirmation.clone()),
        &confirmation,
    );
    let _ = [
        facade::FragmentStageV1::Queue,
        facade::FragmentStageV1::Transcript,
        facade::FragmentStageV1::SpeakerLabels,
        facade::FragmentStageV1::SpeakerFeatures,
        facade::FragmentStageV1::Structuring,
        facade::FragmentStageV1::LabelConfirmation,
    ];
    let _: Result<facade::StagedFragmentV1, facade::FormatError> =
        facade::decode_staged_fragment(&[]);
    let _ = facade::FormatError::InvalidTxIdSlotLength(facade::TxIdSlot::LEN);
}

#[test]
fn v5_event_person_identity_bytes_are_known_or_unknown() {
    let known = person(80);
    let confirmation = facade::LabelConfirmationV1 {
        fragment_id: txid(6),
        interim_txid: txid(7),
        speakers: vec![
            facade::SpeakerLabelV1 {
                speaker: facade::LocalSpeakerLabel::new(1).expect("known speaker label"),
                person_id: Some(known),
            },
            facade::SpeakerLabelV1 {
                speaker: facade::LocalSpeakerLabel::new(2).expect("unknown speaker label"),
                person_id: None,
            },
        ],
    };
    let event = facade::AudioClassificationEventV3::LabelConfirmation(confirmation);
    let mut expected = vec![5, 5];
    expected.extend_from_slice(txid(6).as_bytes());
    expected.extend_from_slice(txid(7).as_bytes());
    expected.extend_from_slice(&[2, 9]);
    expected.extend_from_slice(b"Speaker 1");
    expected.push(1);
    expected.extend_from_slice(known.as_tx_id().as_bytes());
    expected.push(9);
    expected.extend_from_slice(b"Speaker 2");
    expected.push(0);
    assert_eq!(facade::encode_event(&event), Ok(expected.clone()));
    assert_eq!(facade::decode_event(&expected), Ok(event));
}

#[test]
fn speaker_0_3_and_0_2_preserve_event_v5_bytes() {
    let current = current_analysis();
    let previous = previous_analysis();
    assert_eq!(
        postcard::to_allocvec(&current).expect("serialize current analysis"),
        postcard::to_allocvec(&previous).expect("serialize previous analysis")
    );
    let fragment_bytes = [7; 12];
    let current_event =
        facade::encode_event(&facade::AudioClassificationEventV3::TranscriptionComplete(
            facade::TranscriptionCompleteV1 {
                fragment_id: facade::TxId::from_bytes(fragment_bytes),
                analysis: current,
            },
        ))
        .expect("encode current event");
    let previous_event = frozen_event(
        2,
        &PreviousTranscriptionCompleteV1 {
            fragment_id: fragment_bytes,
            analysis: previous,
        },
    );
    assert_eq!(current_event, previous_event);
    assert_eq!(current_event[0], 5);
}

#[test]
fn speaker_0_3_and_0_2_preserve_fragment_v2_bytes() {
    let features = facade::FeatureVector24::default();
    let previous_features = previous::FeatureVector24::default();
    assert_eq!(
        postcard::to_allocvec(&features).expect("serialize current features"),
        postcard::to_allocvec(&previous_features).expect("serialize previous features")
    );
    assert_eq!(
        facade::LocalSpeakerLabel::new(1)
            .expect("current speaker label")
            .to_string(),
        previous::LocalSpeakerLabel::new(1)
            .expect("previous speaker label")
            .to_string()
    );

    let staged = facade::StagedFragmentV1 {
        analysis_txid: txid(20),
        transcript: "hello".into(),
        speakers: vec![facade::StagedSpeakerV1 {
            speaker: facade::LocalSpeakerLabel::new(1).expect("speaker label"),
            language: "English".into(),
            features: features.clone(),
            usable_for_training: true,
        }],
    };
    let staged_bytes = facade::encode_staged_fragment(&staged).expect("encode staged");
    assert_eq!(
        staged_bytes,
        previous_staged_fragment_bytes(&previous_features)
    );
    assert_eq!(staged_bytes[0], 2);
    assert_eq!(facade::decode_staged_fragment(&staged_bytes), Ok(staged));

    let known = person(90);
    let final_value = facade::FinalFragmentV1 {
        analysis_txid: txid(30),
        confirmation_txid: txid(40),
        transcript: "hello".into(),
        speakers: vec![facade::FinalSpeakerV1 {
            speaker: facade::LocalSpeakerLabel::new(1).expect("speaker label"),
            person_id: Some(known),
            language: "English".into(),
            features,
            usable_for_training: false,
        }],
    };
    let final_bytes = facade::encode_final_fragment(&final_value).expect("encode final");
    assert_eq!(
        final_bytes,
        previous_final_fragment_bytes(&previous_features, known)
    );
    assert_eq!(final_bytes[0], 2);
    let identity_offset =
        48 + 24 + final_value.transcript.len() + final_value.speakers[0].speaker.to_string().len();
    assert_eq!(final_bytes[identity_offset], 1);
    assert_eq!(
        &final_bytes[identity_offset + 1..identity_offset + 13],
        known.as_tx_id().as_bytes()
    );
    assert_eq!(
        facade::decode_final_fragment(&final_bytes),
        Ok(final_value.clone())
    );

    let mut unknown = final_value;
    unknown.speakers[0].person_id = None;
    let unknown_bytes = facade::encode_final_fragment(&unknown).expect("encode unknown");
    assert_eq!(unknown_bytes[identity_offset], 0);
    assert_eq!(
        &final_bytes[..identity_offset],
        &unknown_bytes[..identity_offset]
    );
    assert_eq!(
        &final_bytes[identity_offset + 13..],
        &unknown_bytes[identity_offset + 1..]
    );
    assert_eq!(facade::decode_final_fragment(&unknown_bytes), Ok(unknown));

    let slot = facade::TxIdSlot::new(txid(50));
    assert_eq!(facade::TxIdSlot::decode(&slot.encode()), Ok(slot));
    assert_eq!(
        facade::txid_from_path(facade::txid_path(txid(60))),
        Ok(txid(60))
    );
}