kcode-k1-audio-classification-fragment-format 0.2.2

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

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

const FRAGMENT_VERSION: u8 = 2;
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: Option<PersonId>,
    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())?;
        output.push(u8::from(speaker.person_id.is_some()));
        if let Some(person_id) = speaker.person_id {
            output.extend_from_slice(person_id.as_tx_id().as_bytes());
        }
        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_person_id()?,
            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);
    } else {
        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_person_id(&mut self) -> Result<Option<PersonId>, FormatError> {
        if !self.read_boolean()? {
            return Ok(None);
        }
        let bytes = self
            .take(12)?
            .try_into()
            .expect("person ID length was checked");
        Ok(Some(PersonId::from_tx_id(TxId::from_bytes(bytes))))
    }

    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 (input, output) in bytes.chunks_exact(3).zip(encoded.chunks_exact_mut(4)) {
        output[0] = BASE64_ALPHABET[(input[0] >> 2) as usize];
        output[1] = BASE64_ALPHABET[(((input[0] & 3) << 4) | (input[1] >> 4)) as usize];
        output[2] = BASE64_ALPHABET[(((input[1] & 15) << 2) | (input[2] >> 6)) as usize];
        output[3] = BASE64_ALPHABET[(input[2] & 63) as usize];
    }
    encoded
}

fn decode_txid_base64(encoded: [u8; 16]) -> Result<TxId, FormatError> {
    let mut bytes = [0; 12];
    for (input, output) in encoded.chunks_exact(4).zip(bytes.chunks_exact_mut(3)) {
        let first = decode_base64_character(input[0])?;
        let second = decode_base64_character(input[1])?;
        let third = decode_base64_character(input[2])?;
        let fourth = decode_base64_character(input[3])?;
        output[0] = (first << 2) | (second >> 4);
        output[1] = (second << 4) | (third >> 2);
        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;