kcode-k1-audio-classification-fragment-format 0.1.1

use std::path::{Component, Path, PathBuf};

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

const FRAGMENT_VERSION: u8 = 1;
const STAGED_KIND: u8 = 1;
const FINAL_KIND: u8 = 2;
const ANALYSIS_SLOT_OFFSET: usize = 16;
const CONFIRMATION_SLOT_OFFSET: usize = 32;
const BODY_OFFSET: usize = 48;
const BASE64_ALPHABET: &[u8; 64] =
    b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

#[derive(Debug, Clone, PartialEq)]
pub struct StagedSpeakerV1 {
    pub speaker: LocalSpeakerLabel,
    pub language: String,
    pub features: FeatureVector24,
    pub usable_for_training: bool,
}

#[derive(Debug, Clone, PartialEq)]
pub struct StagedFragmentV1 {
    pub analysis_txid: TxId,
    pub transcript: String,
    pub speakers: Vec<StagedSpeakerV1>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct FinalSpeakerV1 {
    pub speaker: LocalSpeakerLabel,
    pub person_id: String,
    pub language: String,
    pub features: FeatureVector24,
    pub usable_for_training: bool,
}

#[derive(Debug, Clone, PartialEq)]
pub struct FinalFragmentV1 {
    pub analysis_txid: TxId,
    pub confirmation_txid: TxId,
    pub transcript: String,
    pub speakers: Vec<FinalSpeakerV1>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TxIdSlot {
    txid: TxId,
}

impl TxIdSlot {
    pub const LEN: usize = 16;
    pub const PADDING_LEN: usize = 4;

    pub const fn new(txid: TxId) -> Self {
        Self { txid }
    }

    pub const fn txid(self) -> TxId {
        self.txid
    }

    pub fn encode(self) -> [u8; Self::LEN] {
        let mut output = [0; Self::LEN];
        output[..12].copy_from_slice(self.txid.as_bytes());
        output
    }

    pub fn decode(bytes: &[u8]) -> Result<Self, FormatError> {
        if bytes.len() != Self::LEN {
            return Err(FormatError::InvalidTxIdSlotLength(bytes.len()));
        }
        if bytes[12..].iter().any(|byte| *byte != 0) {
            return Err(FormatError::NonZeroPadding);
        }
        let mut txid = [0; 12];
        txid.copy_from_slice(&bytes[..12]);
        Ok(Self::new(TxId::from_bytes(txid)))
    }
}

pub fn encode_staged_fragment(value: &StagedFragmentV1) -> Result<Vec<u8>, FormatError> {
    let mut output = encode_fragment_prefix(STAGED_KIND, value.analysis_txid, None);
    append_string(&mut output, &value.transcript)?;
    append_length(&mut output, value.speakers.len())?;
    for speaker in &value.speakers {
        append_string(&mut output, &speaker.speaker.to_string())?;
        append_string(&mut output, &speaker.language)?;
        append_features(&mut output, &speaker.features)?;
        output.push(u8::from(speaker.usable_for_training));
    }
    Ok(output)
}

pub fn decode_staged_fragment(bytes: &[u8]) -> Result<StagedFragmentV1, FormatError> {
    let (analysis_txid, _) = decode_fragment_header(bytes, STAGED_KIND, false)?;
    let mut decoder = BodyDecoder::new(&bytes[BODY_OFFSET..]);
    let transcript = decoder.read_string()?;
    let speaker_count = decoder.read_length()?;
    let mut speakers = Vec::new();
    for _ in 0..speaker_count {
        speakers.push(StagedSpeakerV1 {
            speaker: decoder.read_speaker_label()?,
            language: decoder.read_string()?,
            features: decoder.read_features()?,
            usable_for_training: decoder.read_boolean()?,
        });
    }
    decoder.finish()?;
    Ok(StagedFragmentV1 {
        analysis_txid,
        transcript,
        speakers,
    })
}

pub fn encode_final_fragment(value: &FinalFragmentV1) -> Result<Vec<u8>, FormatError> {
    let mut output = encode_fragment_prefix(
        FINAL_KIND,
        value.analysis_txid,
        Some(value.confirmation_txid),
    );
    append_string(&mut output, &value.transcript)?;
    append_length(&mut output, value.speakers.len())?;
    for speaker in &value.speakers {
        append_string(&mut output, &speaker.speaker.to_string())?;
        append_string(&mut output, &speaker.person_id)?;
        append_string(&mut output, &speaker.language)?;
        append_features(&mut output, &speaker.features)?;
        output.push(u8::from(speaker.usable_for_training));
    }
    Ok(output)
}

pub fn decode_final_fragment(bytes: &[u8]) -> Result<FinalFragmentV1, FormatError> {
    let (analysis_txid, confirmation_txid) = decode_fragment_header(bytes, FINAL_KIND, true)?;
    let mut decoder = BodyDecoder::new(&bytes[BODY_OFFSET..]);
    let transcript = decoder.read_string()?;
    let speaker_count = decoder.read_length()?;
    let mut speakers = Vec::new();
    for _ in 0..speaker_count {
        speakers.push(FinalSpeakerV1 {
            speaker: decoder.read_speaker_label()?,
            person_id: decoder.read_string()?,
            language: decoder.read_string()?,
            features: decoder.read_features()?,
            usable_for_training: decoder.read_boolean()?,
        });
    }
    decoder.finish()?;
    Ok(FinalFragmentV1 {
        analysis_txid,
        confirmation_txid: confirmation_txid
            .expect("final fragment header always contains a confirmation slot"),
        transcript,
        speakers,
    })
}

pub fn txid_path(txid: TxId) -> PathBuf {
    let encoded =
        String::from_utf8(encode_txid_base64(txid).to_vec()).expect("base64 alphabet is ASCII");
    PathBuf::from(&encoded[..1]).join(format!("{}.dat", &encoded[1..]))
}

pub fn txid_from_path(path: impl AsRef<Path>) -> Result<TxId, FormatError> {
    let mut components = path.as_ref().components();
    let shard = normal_utf8_component(components.next())?;
    let filename = normal_utf8_component(components.next())?;
    if components.next().is_some() || shard.len() != 1 {
        return Err(FormatError::InvalidPath);
    }
    let name = filename
        .strip_suffix(".dat")
        .ok_or(FormatError::InvalidPath)?;
    if name.len() != 15 {
        return Err(FormatError::InvalidPath);
    }
    let mut encoded = [0; 16];
    encoded[0] = shard.as_bytes()[0];
    encoded[1..].copy_from_slice(name.as_bytes());
    decode_txid_base64(encoded)
}

fn encode_fragment_prefix(
    kind: u8,
    analysis_txid: TxId,
    confirmation_txid: Option<TxId>,
) -> Vec<u8> {
    let mut output = vec![0; BODY_OFFSET];
    output[0] = FRAGMENT_VERSION;
    output[1] = kind;
    output[ANALYSIS_SLOT_OFFSET..CONFIRMATION_SLOT_OFFSET]
        .copy_from_slice(&TxIdSlot::new(analysis_txid).encode());
    if let Some(txid) = confirmation_txid {
        output[CONFIRMATION_SLOT_OFFSET..BODY_OFFSET]
            .copy_from_slice(&TxIdSlot::new(txid).encode());
    }
    output
}

fn decode_fragment_header(
    bytes: &[u8],
    expected_kind: u8,
    has_confirmation: bool,
) -> Result<(TxId, Option<TxId>), FormatError> {
    if bytes.len() < BODY_OFFSET {
        return Err(FormatError::Truncated);
    }
    if bytes[0] != FRAGMENT_VERSION {
        return Err(FormatError::UnsupportedVersion(bytes[0]));
    }
    if bytes[1] != expected_kind {
        return Err(FormatError::InvalidFragmentKind(bytes[1]));
    }
    let reserved = &bytes[2..ANALYSIS_SLOT_OFFSET];
    if reserved.iter().any(|byte| *byte != 0) {
        return Err(FormatError::NonZeroReserved);
    }
    let analysis_txid =
        TxIdSlot::decode(&bytes[ANALYSIS_SLOT_OFFSET..CONFIRMATION_SLOT_OFFSET])?.txid();
    let confirmation_slot = &bytes[CONFIRMATION_SLOT_OFFSET..BODY_OFFSET];
    let confirmation_txid = if has_confirmation {
        Some(TxIdSlot::decode(confirmation_slot)?.txid())
    } else {
        if confirmation_slot.iter().any(|byte| *byte != 0) {
            return Err(FormatError::NonZeroStagedConfirmation);
        }
        None
    };
    Ok((analysis_txid, confirmation_txid))
}

fn append_length(output: &mut Vec<u8>, length: usize) -> Result<(), FormatError> {
    let length = u64::try_from(length).map_err(|_| FormatError::LengthOverflow)?;
    output.extend_from_slice(&length.to_le_bytes());
    Ok(())
}

fn append_bytes(output: &mut Vec<u8>, bytes: &[u8]) -> Result<(), FormatError> {
    append_length(output, bytes.len())?;
    output.extend_from_slice(bytes);
    Ok(())
}

fn append_string(output: &mut Vec<u8>, value: &str) -> Result<(), FormatError> {
    append_bytes(output, value.as_bytes())
}

fn append_features(output: &mut Vec<u8>, value: &FeatureVector24) -> Result<(), FormatError> {
    let bytes = postcard::to_allocvec(value).map_err(|_| FormatError::InvalidFeatureBody)?;
    append_bytes(output, &bytes)
}

struct BodyDecoder<'a> {
    bytes: &'a [u8],
    position: usize,
}

impl<'a> BodyDecoder<'a> {
    fn new(bytes: &'a [u8]) -> Self {
        Self { bytes, position: 0 }
    }

    fn take(&mut self, length: usize) -> Result<&'a [u8], FormatError> {
        let end = self
            .position
            .checked_add(length)
            .ok_or(FormatError::LengthOverflow)?;
        let value = self
            .bytes
            .get(self.position..end)
            .ok_or(FormatError::Truncated)?;
        self.position = end;
        Ok(value)
    }

    fn read_length(&mut self) -> Result<usize, FormatError> {
        let mut bytes = [0; 8];
        bytes.copy_from_slice(self.take(8)?);
        usize::try_from(u64::from_le_bytes(bytes)).map_err(|_| FormatError::LengthOverflow)
    }

    fn read_bytes(&mut self) -> Result<&'a [u8], FormatError> {
        let length = self.read_length()?;
        self.take(length)
    }

    fn read_string(&mut self) -> Result<String, FormatError> {
        Ok(std::str::from_utf8(self.read_bytes()?)
            .map_err(|_| FormatError::InvalidUtf8)?
            .to_owned())
    }

    fn read_speaker_label(&mut self) -> Result<LocalSpeakerLabel, FormatError> {
        self.read_string()?
            .parse()
            .map_err(|_| FormatError::InvalidSpeakerLabel)
    }

    fn read_features(&mut self) -> Result<FeatureVector24, FormatError> {
        let (value, remaining) = postcard::take_from_bytes(self.read_bytes()?)
            .map_err(|_| FormatError::InvalidFeatureBody)?;
        if !remaining.is_empty() {
            return Err(FormatError::InvalidFeatureBody);
        }
        Ok(value)
    }

    fn read_boolean(&mut self) -> Result<bool, FormatError> {
        match self.take(1)?[0] {
            0 => Ok(false),
            1 => Ok(true),
            value => Err(FormatError::InvalidBoolean(value)),
        }
    }

    fn finish(self) -> Result<(), FormatError> {
        if self.position == self.bytes.len() {
            Ok(())
        } else {
            Err(FormatError::TrailingBytes)
        }
    }
}

fn normal_utf8_component(component: Option<Component<'_>>) -> Result<&str, FormatError> {
    match component {
        Some(Component::Normal(value)) => value.to_str().ok_or(FormatError::InvalidPath),
        _ => Err(FormatError::InvalidPath),
    }
}

fn encode_txid_base64(txid: TxId) -> [u8; 16] {
    let bytes = txid.into_bytes();
    let mut encoded = [0; 16];
    for group in 0..4 {
        let input = group * 3;
        let output = group * 4;
        encoded[output] = BASE64_ALPHABET[(bytes[input] >> 2) as usize];
        encoded[output + 1] =
            BASE64_ALPHABET[(((bytes[input] & 3) << 4) | (bytes[input + 1] >> 4)) as usize];
        encoded[output + 2] =
            BASE64_ALPHABET[(((bytes[input + 1] & 15) << 2) | (bytes[input + 2] >> 6)) as usize];
        encoded[output + 3] = BASE64_ALPHABET[(bytes[input + 2] & 63) as usize];
    }
    encoded
}

fn decode_txid_base64(encoded: [u8; 16]) -> Result<TxId, FormatError> {
    let mut bytes = [0; 12];
    for group in 0..4 {
        let input = group * 4;
        let output = group * 3;
        let first = decode_base64_character(encoded[input])?;
        let second = decode_base64_character(encoded[input + 1])?;
        let third = decode_base64_character(encoded[input + 2])?;
        let fourth = decode_base64_character(encoded[input + 3])?;
        bytes[output] = (first << 2) | (second >> 4);
        bytes[output + 1] = (second << 4) | (third >> 2);
        bytes[output + 2] = (third << 6) | fourth;
    }
    Ok(TxId::from_bytes(bytes))
}

fn decode_base64_character(value: u8) -> Result<u8, FormatError> {
    BASE64_ALPHABET
        .iter()
        .position(|candidate| *candidate == value)
        .and_then(|index| u8::try_from(index).ok())
        .ok_or(FormatError::InvalidPath)
}

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

    fn txid(seed: u8) -> TxId {
        TxId::from_bytes([seed; 12])
    }

    fn staged_fragment() -> StagedFragmentV1 {
        StagedFragmentV1 {
            analysis_txid: txid(20),
            transcript: "hello".into(),
            speakers: vec![StagedSpeakerV1 {
                speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
                language: "English".into(),
                features: FeatureVector24::default(),
                usable_for_training: true,
            }],
        }
    }

    fn final_fragment() -> FinalFragmentV1 {
        FinalFragmentV1 {
            analysis_txid: txid(30),
            confirmation_txid: txid(40),
            transcript: "hello".into(),
            speakers: vec![FinalSpeakerV1 {
                speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
                person_id: "person-1".into(),
                language: "English".into(),
                features: FeatureVector24::default(),
                usable_for_training: false,
            }],
        }
    }

    #[test]
    fn fragments_and_slots_roundtrip_with_exact_offsets() {
        let staged = staged_fragment();
        let staged_bytes = encode_staged_fragment(&staged).expect("encode staged");
        assert_eq!(decode_staged_fragment(&staged_bytes), Ok(staged));

        let final_value = final_fragment();
        let final_bytes = encode_final_fragment(&final_value).expect("encode final");
        assert_eq!(decode_final_fragment(&final_bytes), Ok(final_value));

        assert_eq!(
            TxIdSlot::decode(&[0; 15]),
            Err(FormatError::InvalidTxIdSlotLength(15))
        );
    }

    #[test]
    fn malformed_headers_slots_and_bodies_are_rejected() {
        let staged = staged_fragment();
        let bytes = encode_staged_fragment(&staged).expect("encode staged");
        let transcript = BODY_OFFSET + 8;
        let label = transcript + staged.transcript.len() + 16;
        let feature_length = label + staged.speakers[0].speaker.to_string().len() + 15;
        let feature = feature_length + 8;
        let boolean = feature
            + postcard::to_allocvec(&staged.speakers[0].features)
                .expect("features")
                .len();
        for (index, value, error) in [
            (0, 2, FormatError::UnsupportedVersion(2)),
            (1, 2, FormatError::InvalidFragmentKind(2)),
            (2, 1, FormatError::NonZeroReserved),
            (28, 1, FormatError::NonZeroPadding),
            (32, 1, FormatError::NonZeroStagedConfirmation),
            (transcript, 0xff, FormatError::InvalidUtf8),
            (label, b'x', FormatError::InvalidSpeakerLabel),
            (feature, 2, FormatError::InvalidFeatureBody),
            (boolean, 2, FormatError::InvalidBoolean(2)),
        ] {
            let mut malformed = bytes.clone();
            malformed[index] = value;
            assert_eq!(decode_staged_fragment(&malformed), Err(error));
        }
        assert_eq!(
            decode_staged_fragment(&bytes[..47]),
            Err(FormatError::Truncated)
        );
        let mut trailing = bytes;
        trailing.push(0);
        assert_eq!(
            decode_staged_fragment(&trailing),
            Err(FormatError::TrailingBytes)
        );
    }

    #[test]
    fn every_path_shard_roundtrips_and_noncanonical_paths_fail() {
        for (index, expected) in BASE64_ALPHABET.iter().enumerate() {
            let mut bytes = [0; 12];
            bytes[0] = (index as u8) << 2;
            let txid = TxId::from_bytes(bytes);
            let path = txid_path(txid);
            assert_eq!(path.to_string_lossy().as_bytes()[0], *expected);
            assert_eq!(txid_from_path(&path), Ok(txid));
        }
        for path in [
            "/A/AAAAAAAAAAAAAAA.dat",
            "AA/AAAAAAAAAAAAAAA.dat",
            "A/AAAAAAAAAAAAAA!.dat",
            "A/AAAAAAAAAAAAAAA.dat/extra",
        ] {
            assert_eq!(txid_from_path(path), Err(FormatError::InvalidPath));
        }
    }
}