kcode-speaker-system 0.1.1

use crate::{
    CommitReceipt, Decision, FEATURE_COUNT, FeatureVector, Identification, Key, SampleState,
    SampleStateRequest, SpeakerSystem, SystemError,
};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use std::{error::Error, fmt};

pub const IDENTIFY_KTOOL: &str = "kcode-speaker-system/identify";
pub const SET_SAMPLE_STATE_KTOOL: &str = "kcode-speaker-system/set-sample-state";

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct KtoolSpec {
    pub name: &'static str,
    pub description: &'static str,
    pub input_schema: &'static str,
}

pub const KTOOLS: &[KtoolSpec] = &[
    KtoolSpec {
        name: IDENTIFY_KTOOL,
        description: "Identify one frozen 24-rating speaker profile using the explicitly loaded immutable model and return decision and LLR evidence.",
        input_schema: r#"{"cohortId":"gemini-speaker-24-normalized/1","ratings":[24 integer values in 0..=100]}"#,
    },
    KtoolSpec {
        name: SET_SAMPLE_STATE_KTOOL,
        description: "Change one successor sample to active, confirmed with a speaker ID, or retracted using provenance-bearing event and sample IDs.",
        input_schema: r#"{"eventId":"validated key","sampleId":"store-issued sample ID","reason":"nonblank reason","state":{"status":"active"|"confirmed"|"retracted","speakerId":"required only for confirmed"}}"#,
    },
];

#[derive(Debug)]
pub enum KtoolError {
    UnknownTool,
    InvalidArguments(String),
    ModelUnavailable,
    Execution(SystemError),
    Serialization(String),
}

impl KtoolError {
    pub fn code(&self) -> &'static str {
        match self {
            Self::UnknownTool => "unknown_ktool",
            Self::InvalidArguments(_) => "invalid_arguments",
            Self::ModelUnavailable => "model_unavailable",
            Self::Execution(_) | Self::Serialization(_) => "execution_failed",
        }
    }
}

impl fmt::Display for KtoolError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::UnknownTool => formatter.write_str("unknown_ktool"),
            Self::InvalidArguments(message) => {
                write!(formatter, "invalid_arguments: {message}")
            }
            Self::ModelUnavailable => formatter.write_str("model_unavailable"),
            Self::Execution(error) => write!(formatter, "execution_failed: {error}"),
            Self::Serialization(message) => {
                write!(formatter, "execution_failed: {message}")
            }
        }
    }
}

impl Error for KtoolError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Execution(error) => Some(error),
            Self::UnknownTool
            | Self::InvalidArguments(_)
            | Self::ModelUnavailable
            | Self::Serialization(_) => None,
        }
    }
}

pub fn execute(system: &SpeakerSystem, name: &str, arguments: &str) -> Result<String, KtoolError> {
    match name {
        IDENTIFY_KTOOL => execute_identify(system, arguments),
        SET_SAMPLE_STATE_KTOOL => execute_set_sample_state(system, arguments),
        _ => Err(KtoolError::UnknownTool),
    }
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct IdentifyArguments {
    cohort_id: Key,
    ratings: [u8; FEATURE_COUNT],
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct IdentifyResponse {
    decision: DecisionResponse,
    best: CandidateResponse,
    runner_up: Option<CandidateResponse>,
    absolute_pass: bool,
    margin_pass: bool,
}

#[derive(Serialize)]
#[serde(tag = "status", rename_all = "snake_case")]
enum DecisionResponse {
    Known {
        #[serde(rename = "speakerId")]
        speaker_id: Key,
    },
    Unknown,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct CandidateResponse {
    speaker_id: Key,
    llr: f64,
}

fn execute_identify(system: &SpeakerSystem, arguments: &str) -> Result<String, KtoolError> {
    let arguments: IdentifyArguments = decode_arguments(arguments)?;
    if &arguments.cohort_id != crate::cohort_id() {
        return Err(KtoolError::InvalidArguments(
            "cohortId must equal gemini-speaker-24-normalized/1".into(),
        ));
    }
    let features = FeatureVector::new(arguments.ratings).map_err(|error| {
        KtoolError::InvalidArguments(format!("ratings must be integers in 0..=100: {error}"))
    })?;
    let identification = system
        .identify(&arguments.cohort_id, &features)
        .map_err(map_identify_error)?;
    encode_response(identify_response(identification))
}

fn identify_response(identification: Identification) -> IdentifyResponse {
    let decision = match identification.decision {
        Decision::Known { speaker_id } => DecisionResponse::Known { speaker_id },
        Decision::Unknown => DecisionResponse::Unknown,
    };
    IdentifyResponse {
        decision,
        best: CandidateResponse {
            speaker_id: identification.best.speaker_id,
            llr: identification.best.llr,
        },
        runner_up: identification.runner_up.map(|candidate| CandidateResponse {
            speaker_id: candidate.speaker_id,
            llr: candidate.llr,
        }),
        absolute_pass: identification.absolute_pass,
        margin_pass: identification.margin_pass,
    }
}

fn map_identify_error(error: SystemError) -> KtoolError {
    match error {
        SystemError::ModelUnavailable => KtoolError::ModelUnavailable,
        other => KtoolError::Execution(other),
    }
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct SetSampleStateArguments {
    event_id: Key,
    sample_id: Key,
    reason: String,
    state: StateArgument,
}

#[derive(Deserialize)]
#[serde(untagged)]
enum StateArgument {
    Active(ActiveState),
    Confirmed(ConfirmedState),
    Retracted(RetractedState),
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ActiveState {
    status: ActiveStatus,
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ConfirmedState {
    status: ConfirmedStatus,
    #[serde(rename = "speakerId")]
    speaker_id: Key,
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct RetractedState {
    status: RetractedStatus,
}

#[derive(Deserialize)]
enum ActiveStatus {
    #[serde(rename = "active")]
    Active,
}

#[derive(Deserialize)]
enum ConfirmedStatus {
    #[serde(rename = "confirmed")]
    Confirmed,
}

#[derive(Deserialize)]
enum RetractedStatus {
    #[serde(rename = "retracted")]
    Retracted,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct StateResponse {
    revision: u64,
    applied: u64,
}

fn execute_set_sample_state(system: &SpeakerSystem, arguments: &str) -> Result<String, KtoolError> {
    let arguments: SetSampleStateArguments = decode_arguments(arguments)?;
    if arguments.reason.trim().is_empty() {
        return Err(KtoolError::InvalidArguments(
            "reason must be nonblank".into(),
        ));
    }

    let state = match arguments.state {
        StateArgument::Active(state) => {
            let ActiveStatus::Active = state.status;
            SampleState::Unlabeled
        }
        StateArgument::Confirmed(state) => {
            let ConfirmedStatus::Confirmed = state.status;
            SampleState::Confirmed {
                speaker_id: state.speaker_id,
            }
        }
        StateArgument::Retracted(state) => {
            let RetractedStatus::Retracted = state.status;
            SampleState::Retracted
        }
    };

    let receipt = system
        .change_sample_state(SampleStateRequest {
            event_id: arguments.event_id,
            sample_id: arguments.sample_id,
            state,
            reason: arguments.reason,
        })
        .map_err(KtoolError::Execution)?;
    encode_receipt(receipt)
}

fn decode_arguments<T: DeserializeOwned>(arguments: &str) -> Result<T, KtoolError> {
    serde_json::from_str(arguments).map_err(|error| KtoolError::InvalidArguments(error.to_string()))
}

fn encode_receipt(receipt: CommitReceipt) -> Result<String, KtoolError> {
    encode_response(StateResponse {
        revision: receipt.revision,
        applied: receipt.applied,
    })
}

fn encode_response<T: Serialize>(response: T) -> Result<String, KtoolError> {
    serde_json::to_string(&response).map_err(|error| KtoolError::Serialization(error.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        AttemptSelectionRequest, NormalizedAttempt, ObjectId, RecordingKind, SegmentBinding,
        SourceRegistration, cohort_id, open,
    };
    use serde_json::{Value, json};
    use std::{
        fs,
        path::PathBuf,
        sync::atomic::{AtomicU64, Ordering},
    };

    static NEXT_PATH: AtomicU64 = AtomicU64::new(0);

    fn key(value: &str) -> Key {
        Key::parse(value).unwrap()
    }

    fn object(value: &str) -> ObjectId {
        ObjectId::parse(value).unwrap()
    }

    fn path() -> PathBuf {
        std::env::temp_dir().join(format!(
            "kcode-speaker-system-ktool-{}-{}.db",
            std::process::id(),
            NEXT_PATH.fetch_add(1, Ordering::Relaxed)
        ))
    }

    fn ratings() -> Vec<u8> {
        (10..34).collect()
    }

    fn registered_sample(system: &SpeakerSystem) -> Key {
        system
            .register_source(SourceRegistration {
                source_object: object("SOURCE31"),
                source_duration_ms: 100_000,
                group_id: key("group/source/31"),
                recording_kind: RecordingKind::VoiceNote,
                segments: vec![SegmentBinding {
                    event_id: key("event/register/31"),
                    clip_object: object("CLIP0031"),
                }],
            })
            .unwrap();

        let feature_list = ratings()
            .into_iter()
            .map(|rating| rating.to_string())
            .collect::<Vec<_>>()
            .join(",");
        let normalized_response = format!(
            r#"{{"status":"scored","speakers":[{{"speakerOrdinal":0,"primaryLanguage":"en-US","closestDialect":"General American English","usableSpeechMs":20000,"features":[{feature_list}]}}],"additionalSpeakers":[]}}"#
        );
        system
            .record_normalized_attempt(NormalizedAttempt {
                event_id: key("event/attempt/31"),
                attempt_id: key("attempt/31"),
                cohort_id: cohort_id().clone(),
                source_object: object("SOURCE31"),
                clip_object: object("CLIP0031"),
                segment_ordinal: 0,
                provider_result_object: object("RESULT31"),
                recording_quality: Some(90),
                normalized_response: &normalized_response,
            })
            .unwrap();
        system
            .select_attempt(AttemptSelectionRequest {
                event_id: key("event/select/31"),
                attempt_id: key("attempt/31"),
                cohort_id: cohort_id().clone(),
                clip_object: object("CLIP0031"),
                reason: "selected complete successor sample".into(),
            })
            .unwrap();

        system
            .attempt(&key("attempt/31"))
            .unwrap()
            .unwrap()
            .sample_ids[0]
            .clone()
    }

    #[test]
    fn identify_json_is_strict_and_model_unavailable_is_explicit() {
        let database = path();
        let system = open(&database).unwrap();
        let valid = json!({
            "cohortId": "gemini-speaker-24-normalized/1",
            "ratings": ratings(),
        });

        let error = execute(&system, IDENTIFY_KTOOL, &valid.to_string()).unwrap_err();
        assert!(matches!(error, KtoolError::ModelUnavailable));
        assert_eq!(error.code(), "model_unavailable");
        assert_eq!(error.to_string(), "model_unavailable");

        let unknown = json!({
            "cohortId": "gemini-speaker-24-normalized/1",
            "ratings": ratings(),
            "speaker": "legacy",
        });
        assert!(matches!(
            execute(&system, IDENTIFY_KTOOL, &unknown.to_string()),
            Err(KtoolError::InvalidArguments(_))
        ));

        let legacy = json!({
            "cohort": "legacy",
            "features": ratings(),
        });
        assert!(matches!(
            execute(&system, IDENTIFY_KTOOL, &legacy.to_string()),
            Err(KtoolError::InvalidArguments(_))
        ));

        let mut short = ratings();
        short.pop();
        let wrong_count = json!({
            "cohortId": "gemini-speaker-24-normalized/1",
            "ratings": short,
        });
        assert!(matches!(
            execute(&system, IDENTIFY_KTOOL, &wrong_count.to_string()),
            Err(KtoolError::InvalidArguments(_))
        ));

        let mut out_of_range = ratings();
        out_of_range[0] = 101;
        let wrong_range = json!({
            "cohortId": "gemini-speaker-24-normalized/1",
            "ratings": out_of_range,
        });
        assert!(matches!(
            execute(&system, IDENTIFY_KTOOL, &wrong_range.to_string()),
            Err(KtoolError::InvalidArguments(_))
        ));

        let wrong_cohort = json!({
            "cohortId": "gemini-speaker-24-freeform/1",
            "ratings": ratings(),
        });
        assert!(matches!(
            execute(&system, IDENTIFY_KTOOL, &wrong_cohort.to_string()),
            Err(KtoolError::InvalidArguments(_))
        ));

        drop(system);
        fs::remove_file(database).unwrap();
    }

    #[test]
    fn state_json_rejects_legacy_before_mutation_and_changes_successor_state() {
        let database = path();
        let system = open(&database).unwrap();
        let sample_id = registered_sample(&system);
        let event_id = "event/state/ktool/31";

        let legacy = json!({
            "eventId": event_id,
            "sampleId": sample_id,
            "reason": "legacy operation must fail",
            "action": "train",
            "speakerId": "speaker/alice",
        });
        assert!(matches!(
            execute(&system, SET_SAMPLE_STATE_KTOOL, &legacy.to_string()),
            Err(KtoolError::InvalidArguments(_))
        ));

        let confirmed = json!({
            "eventId": event_id,
            "sampleId": sample_id,
            "reason": "confirmed from successor sample provenance",
            "state": {
                "status": "confirmed",
                "speakerId": "speaker/alice",
            },
        });
        let response = execute(&system, SET_SAMPLE_STATE_KTOOL, &confirmed.to_string()).unwrap();
        let response: Value = serde_json::from_str(&response).unwrap();
        assert_eq!(response["applied"], 1);
        assert_eq!(
            kcode_speaker_dataset::active_rows(&system.dataset(cohort_id()).unwrap()).len(),
            1
        );

        let active = json!({
            "eventId": "event/state/ktool/32",
            "sampleId": sample_id,
            "reason": "return sample to active unconfirmed review",
            "state": {
                "status": "active",
            },
        });
        execute(&system, SET_SAMPLE_STATE_KTOOL, &active.to_string()).unwrap();
        assert!(matches!(
            system.dataset(cohort_id()),
            Err(crate::SystemError::Dataset(
                crate::DatasetError::EmptyActive
            ))
        ));

        let retracted = json!({
            "eventId": "event/state/ktool/33",
            "sampleId": sample_id,
            "reason": "retract successor sample",
            "state": {
                "status": "retracted",
            },
        });
        execute(&system, SET_SAMPLE_STATE_KTOOL, &retracted.to_string()).unwrap();
        assert!(matches!(
            system.dataset(cohort_id()),
            Err(crate::SystemError::Dataset(
                crate::DatasetError::EmptyActive
            ))
        ));

        let speaker_on_active = json!({
            "eventId": "event/state/ktool/34",
            "sampleId": sample_id,
            "reason": "unknown nested field must fail",
            "state": {
                "status": "active",
                "speakerId": "speaker/alice",
            },
        });
        assert!(matches!(
            execute(
                &system,
                SET_SAMPLE_STATE_KTOOL,
                &speaker_on_active.to_string()
            ),
            Err(KtoolError::InvalidArguments(_))
        ));

        drop(system);
        fs::remove_file(database).unwrap();
    }

    #[test]
    fn only_successor_names_dispatch() {
        let database = path();
        let system = open(&database).unwrap();

        assert_eq!(KTOOLS.len(), 2);
        assert_eq!(KTOOLS[0].name, IDENTIFY_KTOOL);
        assert_eq!(KTOOLS[1].name, SET_SAMPLE_STATE_KTOOL);
        assert!(matches!(
            execute(&system, "kcode-speech-classifier/train", "{}"),
            Err(KtoolError::UnknownTool)
        ));

        drop(system);
        fs::remove_file(database).unwrap();
    }
}