kcode-k1-audio-classification-projection-fold 0.1.0

use kcode_k1_audio_classification_event_format::{
    AttemptFinalAnalysisV1, AudioClassificationEventV3 as V5, AudioClassificationEventV6 as V6,
    decode_event, decode_event_v6, decode_final_event_v7,
};
use kcode_k1_audio_classification_projection_artifacts::AttemptArtifacts;
use kcode_k1_audio_classification_projection_state as state;
use kcode_k1_audio_classification_projection_v5_fold::fold_v5 as fold_decoded_v5;
use kcode_k1_audio_classification_projection_v6_fold::fold_v6 as fold_decoded_v6;

pub use kcode_k1_audio_classification_projection_artifacts::{
    LocalSpeakerLabel, SpeakerFeatureEvidence, TxId,
};
pub use kcode_k1_audio_classification_projection_state::{
    ExecutedAnalysis, FragmentId, FragmentStageV1, FragmentStatus, LlmJobState, LlmJobStatus,
    OverallState, SpeakerLabelV1, StageState, StageStatus,
};
pub use kcode_k1_audio_classification_projection_v5_fold::ProjectionEffect;

#[derive(Clone)]
pub struct FragmentProjection {
    status: FragmentStatus,
    artifacts: AttemptArtifacts,
}

#[derive(Clone, Debug, PartialEq)]
pub struct ResumePlan {
    pub fragment_id: FragmentId,
    pub transcript: Option<String>,
    pub speaker_labels: Option<Vec<LocalSpeakerLabel>>,
    pub speaker_features: Vec<SpeakerFeatureEvidence>,
    pub final_analysis: Option<ExecutedAnalysis>,
    pub active_attempt: Option<TxId>,
    pub interrupted: bool,
}

#[derive(Clone)]
pub struct FoldResult {
    pub fragment_id: FragmentId,
    pub effect: ProjectionEffect,
    pub replacement: Option<FragmentProjection>,
}

impl FragmentProjection {
    pub fn fragment_id(&self) -> FragmentId {
        self.artifacts.fragment_id()
    }

    pub fn visible_status(&self) -> FragmentStatus {
        let mut status = self.status.clone();
        overlay_stages(&mut status, &self.artifacts);
        status.state = projected_state(&status, &self.artifacts);
        status
    }

    pub fn resume_plan(&self) -> ResumePlan {
        let interrupted = !self.artifacts.complete() && self.status.state == OverallState::Running;
        ResumePlan {
            fragment_id: self.fragment_id(),
            transcript: self.artifacts.transcript().map(str::to_owned),
            speaker_labels: self.artifacts.speaker_labels().map(<[_]>::to_vec),
            speaker_features: self.artifacts.speaker_features().to_vec(),
            final_analysis: self.artifacts.final_analysis().cloned(),
            active_attempt: self.artifacts.active_attempt().filter(|_| interrupted),
            interrupted,
        }
    }

    pub fn needs_work(&self) -> bool {
        let state = self.visible_status().state;
        !self.artifacts.complete() && matches!(state, OverallState::Queued | OverallState::Running)
    }

    pub fn validate_labels(&self, labels: &[SpeakerLabelV1]) -> Result<TxId, String> {
        let mut status = self.status.clone();
        if status.analysis.is_some() && status.interim_txid.is_some() {
            status.state = OverallState::Completed;
        }
        state::validate_labels(&status, labels)
    }

    #[cfg(feature = "testkit")]
    pub fn inject_errors(&mut self, errors: Vec<String>) {
        for error in errors {
            state::append_error(&mut self.status, error);
        }
    }
}

pub fn fold(
    current: Option<&FragmentProjection>,
    callback_txid: TxId,
    payload: &[u8],
) -> Result<FoldResult, String> {
    match payload.first() {
        Some(6) => fold_v6(
            current,
            callback_txid,
            &decode_event_v6(payload).map_err(message)?,
        ),
        Some(7) => fold_v7(
            current,
            callback_txid,
            &decode_final_event_v7(payload).map_err(message)?,
        ),
        _ => fold_v5(
            current,
            callback_txid,
            &decode_event(payload).map_err(message)?,
        ),
    }
}

fn fold_v5(
    current: Option<&FragmentProjection>,
    callback: TxId,
    event: &V5,
) -> Result<FoldResult, String> {
    let update = fold_decoded_v5(
        current.map(|value| (&value.status, &value.artifacts)),
        callback,
        event,
    )?;
    let replacement = update
        .replacement
        .map(|value| projection(value.status, value.artifacts));
    Ok(FoldResult {
        fragment_id: update.fragment_id,
        effect: update.effect,
        replacement,
    })
}

fn fold_v6(
    current: Option<&FragmentProjection>,
    callback: TxId,
    event: &V6,
) -> Result<FoldResult, String> {
    let fragment = v6_fragment(event);
    let current = current.ok_or_else(|| "event references an unknown fragment".to_string())?;
    require(fragment == current.fragment_id(), "V6 fragment mismatch")?;
    let update = fold_decoded_v6(&current.status, &current.artifacts, callback, event)?;
    let replacement = update
        .replacement
        .map(|value| projection(value.status, value.artifacts));
    Ok(none_effect(fragment, replacement))
}

fn fold_v7(
    current: Option<&FragmentProjection>,
    callback: TxId,
    event: &AttemptFinalAnalysisV1,
) -> Result<FoldResult, String> {
    let fragment = event.fragment_id;
    let current = current.ok_or_else(|| "event references an unknown fragment".to_string())?;
    require(fragment == current.fragment_id(), "V7 fragment mismatch")?;
    if current.status.state != OverallState::Running
        || current.artifacts.active_attempt() != Some(event.attempt_txid)
    {
        return Ok(none_effect(fragment, None));
    }
    let update = current.artifacts.fold_v7(callback, event)?;
    let Some(artifacts) = update.replacement else {
        return Ok(none_effect(fragment, None));
    };
    let analysis = artifacts
        .final_analysis()
        .cloned()
        .ok_or("final analysis is absent after V7")?;
    let mut status = current.status.clone();
    status.analysis = Some(analysis);
    status.interim_txid = Some(callback);
    status.transcript.state = StageState::Succeeded;
    status.speaker_labels.state = StageState::Succeeded;
    status.speaker_features.state = StageState::Succeeded;
    status.structuring.state = StageState::Succeeded;
    status.state = projected_state(&status, &artifacts);
    Ok(none_effect(fragment, Some(projection(status, artifacts))))
}

fn overlay_stages(status: &mut FragmentStatus, artifacts: &AttemptArtifacts) {
    let terminal = matches!(
        status.state,
        OverallState::Completed | OverallState::Confirmed
    );
    for (stage, artifact_stage) in [
        (&mut status.transcript, FragmentStageV1::Transcript),
        (&mut status.speaker_labels, FragmentStageV1::SpeakerLabels),
        (
            &mut status.speaker_features,
            FragmentStageV1::SpeakerFeatures,
        ),
        (&mut status.structuring, FragmentStageV1::Structuring),
    ] {
        if artifacts.stage_present(artifact_stage) {
            stage.state = StageState::Succeeded;
        } else if terminal {
            stage.state = StageState::Pending;
        }
    }
}

fn projected_state(status: &FragmentStatus, artifacts: &AttemptArtifacts) -> OverallState {
    let confirmed = artifacts.complete()
        && status.label_confirmation.state == StageState::Succeeded
        && status.final_transcript.is_some();
    match status.state {
        OverallState::Discarded => OverallState::Discarded,
        _ if confirmed => OverallState::Confirmed,
        _ if artifacts.complete() => OverallState::Completed,
        OverallState::Failed => OverallState::Failed,
        OverallState::Running => OverallState::Running,
        _ => OverallState::Queued,
    }
}

fn v6_fragment(event: &V6) -> FragmentId {
    match event {
        V6::AttemptStarted(value) => value.fragment_id,
        V6::Progress(value) => value.fragment_id,
        V6::GeminiTranscript(value) => value.fragment_id,
        V6::TerraSpeakerLabels(value) => value.fragment_id,
        V6::GeminiFeatureBundle(value) => value.fragment_id,
        V6::Failed(value) => value.fragment_id,
    }
}

fn projection(status: FragmentStatus, artifacts: AttemptArtifacts) -> FragmentProjection {
    FragmentProjection { status, artifacts }
}

fn message(error: impl ToString) -> String {
    error.to_string()
}

fn require(condition: bool, message: &str) -> Result<(), String> {
    condition.then_some(()).ok_or_else(|| message.to_string())
}

fn none_effect(fragment_id: FragmentId, replacement: Option<FragmentProjection>) -> FoldResult {
    FoldResult {
        fragment_id,
        effect: ProjectionEffect::None,
        replacement,
    }
}

#[cfg(test)]
mod tests;