kcode-k1-audio-classification-event-format 0.2.1

use kcode_k1_person_types::PersonId;
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 = 5;
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,
    #[serde(with = "person_id_option_serde")]
    pub person_id: Option<PersonId>,
}

#[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)?))
    }
}

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

    pub fn serialize<S: Serializer>(
        value: &Option<PersonId>,
        serializer: S,
    ) -> Result<S::Ok, S::Error> {
        value
            .map(|person_id| person_id.as_tx_id().into_bytes())
            .serialize(serializer)
    }

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