internal_tools.rs

//! Registry and dispatch for Muninn builtin pipeline steps.
//!
//! Maps canonical step command names (STT providers and `refine`) onto
//! in-process executors and `__internal_step` subprocess entry points. Pipeline
//! configuration rewrites builtin commands to envelope-json I/O before the
//! [`muninn::PipelineRunner`] runs.

use std::process::ExitCode;

use anyhow::Result;
use async_trait::async_trait;
use muninn::config::{PipelineStepConfig, StepIoMode};
use muninn::{
    InProcessStepError, InProcessStepExecutor, MuninnEnvelopeV1, ResolvedBuiltinStepConfig,
    StepFailureKind, TranscriptionProvider,
};

use crate::{
    refine, stt_apple_speech_tool, stt_deepgram_tool, stt_google_tool, stt_openai_tool,
    stt_whisper_cpp_tool, stt_xai_tool,
};

const INTERNAL_STEP_MARKER: &str = "__internal_step";

/// High-level category for a builtin pipeline step.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BuiltinStepKind {
    /// Speech-to-text provider step.
    Transcription,
    /// Post-transcription transform such as refine.
    Transform,
}

/// Canonical builtin steps recognized by command name.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BuiltinStep {
    /// Apple Speech on-device STT.
    SttAppleSpeech,
    /// Local whisper.cpp STT.
    SttWhisperCpp,
    /// Deepgram cloud STT.
    SttDeepgram,
    /// OpenAI cloud STT.
    SttOpenAi,
    /// Google cloud STT.
    SttGoogle,
    /// SpaceXAI / xAI cloud STT.
    SttXai,
    /// LLM transcript refinement.
    Refine,
}

impl BuiltinStep {
    /// Pipeline `cmd` string for this builtin step.
    pub const fn canonical_name(self) -> &'static str {
        match self {
            Self::SttAppleSpeech => "stt_apple_speech",
            Self::SttWhisperCpp => "stt_whisper_cpp",
            Self::SttDeepgram => "stt_deepgram",
            Self::SttOpenAi => "stt_openai",
            Self::SttGoogle => "stt_google",
            Self::SttXai => "stt_xai",
            Self::Refine => "refine",
        }
    }

    /// Return whether this step is transcription or transform.
    pub const fn kind(self) -> BuiltinStepKind {
        match self {
            Self::SttAppleSpeech
            | Self::SttWhisperCpp
            | Self::SttDeepgram
            | Self::SttOpenAi
            | Self::SttGoogle
            | Self::SttXai => BuiltinStepKind::Transcription,
            Self::Refine => BuiltinStepKind::Transform,
        }
    }

    /// True when this step is an STT provider rather than a transform.
    pub const fn is_transcription(self) -> bool {
        matches!(self.kind(), BuiltinStepKind::Transcription)
    }

    /// Dispatch subprocess `__internal_step` execution for this builtin.
    pub fn run_as_internal_tool(self) -> ExitCode {
        match self {
            Self::SttAppleSpeech => stt_apple_speech_tool::run_as_internal_tool(),
            Self::SttWhisperCpp => stt_whisper_cpp_tool::run_as_internal_tool(),
            Self::SttDeepgram => stt_deepgram_tool::run_as_internal_tool(),
            Self::SttOpenAi => stt_openai_tool::run_as_internal_tool(),
            Self::SttGoogle => stt_google_tool::run_as_internal_tool(),
            Self::SttXai => stt_xai_tool::run_as_internal_tool(),
            Self::Refine => refine::run_as_internal_tool(),
        }
    }

    async fn execute_in_process(
        self,
        input: &MuninnEnvelopeV1,
        config: &ResolvedBuiltinStepConfig,
    ) -> Result<MuninnEnvelopeV1, InProcessStepError> {
        match self {
            Self::SttAppleSpeech => stt_apple_speech_tool::process_input_in_process(input, config)
                .await
                .map_err(map_internal_tool_error),
            Self::SttWhisperCpp => stt_whisper_cpp_tool::process_input_in_process(input, config)
                .await
                .map_err(map_internal_tool_error),
            Self::SttDeepgram => stt_deepgram_tool::process_input_in_process(input, config)
                .await
                .map_err(map_internal_tool_error),
            Self::SttOpenAi => stt_openai_tool::process_input_in_process(input, config)
                .await
                .map_err(map_internal_tool_error),
            Self::SttGoogle => stt_google_tool::process_input_in_process(input, config)
                .await
                .map_err(map_internal_tool_error),
            Self::SttXai => stt_xai_tool::process_input_in_process(input, config)
                .await
                .map_err(map_internal_tool_error),
            Self::Refine => refine::process_input_in_process(input, config)
                .await
                .map_err(map_internal_tool_error),
        }
    }
}

/// [`InProcessStepExecutor`] that runs registered builtin steps in-process.
#[derive(Debug, Clone)]
pub struct BuiltinStepExecutor {
    config: ResolvedBuiltinStepConfig,
}

impl BuiltinStepExecutor {
    /// Build an executor with resolved builtin-step configuration.
    pub fn new(config: ResolvedBuiltinStepConfig) -> Self {
        Self { config }
    }
}

#[async_trait]
impl InProcessStepExecutor for BuiltinStepExecutor {
    async fn try_execute(
        &self,
        step: &PipelineStepConfig,
        input: &MuninnEnvelopeV1,
    ) -> Option<Result<MuninnEnvelopeV1, InProcessStepError>> {
        let builtin = lookup_builtin_step(&step.cmd)?;
        Some(builtin.execute_in_process(input, &self.config).await)
    }
}

/// Handle `muninn __internal_step <name>` CLI dispatch when args match.
///
/// Returns `None` when `args` are not an internal-step invocation.
pub fn maybe_handle_internal_step(args: &[String]) -> Option<ExitCode> {
    if args.get(1).map(String::as_str) != Some(INTERNAL_STEP_MARKER) {
        return None;
    }

    let Some(step_name) = args.get(2).map(String::as_str) else {
        eprintln!("muninn internal step failed: missing step name after {INTERNAL_STEP_MARKER}");
        return Some(ExitCode::FAILURE);
    };

    let Some(tool) = lookup_builtin_step(step_name) else {
        eprintln!("muninn internal step failed: unknown internal step '{step_name}'");
        return Some(ExitCode::FAILURE);
    };

    Some(tool.run_as_internal_tool())
}

/// Normalize a builtin `step.cmd` to its canonical name and envelope-json I/O.
///
/// Returns `Ok(true)` when the command was rewritten, `Ok(false)` for external
/// commands left unchanged.
pub fn rewrite_internal_tool_step(step: &mut PipelineStepConfig) -> Result<bool> {
    let Some(tool) = lookup_builtin_step(&step.cmd) else {
        return Ok(false);
    };

    step.cmd = tool.canonical_name().to_string();
    step.io_mode = StepIoMode::EnvelopeJson;
    Ok(true)
}

/// True when `step.cmd` names a builtin STT provider.
pub fn is_transcription_step(step: &PipelineStepConfig) -> bool {
    lookup_builtin_step(&step.cmd).is_some_and(BuiltinStep::is_transcription)
}

/// Resolve a pipeline command string to a [`BuiltinStep`] when recognized.
pub fn lookup_builtin_step(raw: &str) -> Option<BuiltinStep> {
    if let Some(provider) = TranscriptionProvider::lookup_step_name(raw) {
        return Some(match provider {
            TranscriptionProvider::AppleSpeech => BuiltinStep::SttAppleSpeech,
            TranscriptionProvider::WhisperCpp => BuiltinStep::SttWhisperCpp,
            TranscriptionProvider::Deepgram => BuiltinStep::SttDeepgram,
            TranscriptionProvider::OpenAi => BuiltinStep::SttOpenAi,
            TranscriptionProvider::Google => BuiltinStep::SttGoogle,
            TranscriptionProvider::Xai => BuiltinStep::SttXai,
        });
    }

    match raw {
        "refine" => Some(BuiltinStep::Refine),
        _ => None,
    }
}

fn map_internal_tool_error(error: impl InternalToolError) -> InProcessStepError {
    InProcessStepError {
        kind: StepFailureKind::NonZeroExit,
        message: error.message().to_string(),
        stderr: error.to_stderr_json(),
        exit_status: Some(1),
    }
}

trait InternalToolError {
    fn message(&self) -> &str;
    fn to_stderr_json(&self) -> String;
}

impl InternalToolError for stt_openai_tool::CliError {
    fn message(&self) -> &str {
        stt_openai_tool::CliError::message(self)
    }

    fn to_stderr_json(&self) -> String {
        stt_openai_tool::CliError::to_stderr_json(self)
    }
}

impl InternalToolError for stt_google_tool::CliError {
    fn message(&self) -> &str {
        stt_google_tool::CliError::message(self)
    }

    fn to_stderr_json(&self) -> String {
        stt_google_tool::CliError::to_stderr_json(self)
    }
}

impl InternalToolError for stt_apple_speech_tool::CliError {
    fn message(&self) -> &str {
        stt_apple_speech_tool::CliError::message(self)
    }

    fn to_stderr_json(&self) -> String {
        stt_apple_speech_tool::CliError::to_stderr_json(self)
    }
}

impl InternalToolError for stt_deepgram_tool::CliError {
    fn message(&self) -> &str {
        stt_deepgram_tool::CliError::message(self)
    }

    fn to_stderr_json(&self) -> String {
        stt_deepgram_tool::CliError::to_stderr_json(self)
    }
}

impl InternalToolError for stt_whisper_cpp_tool::CliError {
    fn message(&self) -> &str {
        stt_whisper_cpp_tool::CliError::message(self)
    }

    fn to_stderr_json(&self) -> String {
        stt_whisper_cpp_tool::CliError::to_stderr_json(self)
    }
}

impl InternalToolError for stt_xai_tool::CliError {
    fn message(&self) -> &str {
        stt_xai_tool::CliError::message(self)
    }

    fn to_stderr_json(&self) -> String {
        stt_xai_tool::CliError::to_stderr_json(self)
    }
}

impl InternalToolError for refine::CliError {
    fn message(&self) -> &str {
        refine::CliError::message(self)
    }

    fn to_stderr_json(&self) -> String {
        refine::CliError::to_stderr_json(self)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use muninn::config::{ConfigValidationError, OnErrorPolicy, PipelineStepConfig, StepIoMode};
    use std::path::PathBuf;

    #[test]
    fn normalizes_internal_tool_step_to_canonical_builtin_name() {
        let mut step = PipelineStepConfig {
            id: "refine".to_string(),
            cmd: "refine".to_string(),
            args: vec!["--example".to_string()],
            io_mode: StepIoMode::Auto,
            timeout_ms: 100,
            on_error: OnErrorPolicy::Continue,
        };

        let rewritten = rewrite_internal_tool_step(&mut step).expect("rewrite should succeed");

        assert!(rewritten);
        assert_eq!(step.cmd, "refine");
        assert_eq!(step.args, vec!["--example"]);
        assert_eq!(step.io_mode, StepIoMode::EnvelopeJson);
    }

    #[test]
    fn leaves_external_command_unchanged() {
        let mut step = PipelineStepConfig {
            id: "uppercase".to_string(),
            cmd: "/opt/homebrew/bin/jq".to_string(),
            args: vec!["-c".to_string()],
            io_mode: StepIoMode::Auto,
            timeout_ms: 100,
            on_error: OnErrorPolicy::Continue,
        };

        let rewritten = rewrite_internal_tool_step(&mut step).expect("rewrite should succeed");

        assert!(!rewritten);
        assert_eq!(step.cmd, "/opt/homebrew/bin/jq");
        assert_eq!(step.args, vec!["-c"]);
        assert_eq!(step.io_mode, StepIoMode::Auto);
    }

    #[test]
    fn lookup_builtin_step_accepts_only_canonical_builtin_names() {
        assert_eq!(
            lookup_builtin_step("stt_apple_speech").map(BuiltinStep::canonical_name),
            Some("stt_apple_speech")
        );
        assert_eq!(
            lookup_builtin_step("stt_whisper_cpp").map(BuiltinStep::canonical_name),
            Some("stt_whisper_cpp")
        );
        assert_eq!(
            lookup_builtin_step("stt_deepgram").map(BuiltinStep::canonical_name),
            Some("stt_deepgram")
        );
        assert_eq!(
            lookup_builtin_step("stt_openai").map(BuiltinStep::canonical_name),
            Some("stt_openai")
        );
        assert_eq!(
            lookup_builtin_step("stt_google").map(BuiltinStep::canonical_name),
            Some("stt_google")
        );
        assert_eq!(
            lookup_builtin_step("stt_xai").map(BuiltinStep::canonical_name),
            Some("stt_xai")
        );
        assert_eq!(
            lookup_builtin_step("refine").map(BuiltinStep::canonical_name),
            Some("refine")
        );
        assert_eq!(lookup_builtin_step("muninn-stt-openai"), None);
        assert_eq!(lookup_builtin_step("muninn-stt-google"), None);
        assert_eq!(lookup_builtin_step("muninn-refine"), None);
    }

    #[test]
    fn classifies_transcription_steps_from_registry() {
        assert!(BuiltinStep::SttAppleSpeech.is_transcription());
        assert!(BuiltinStep::SttWhisperCpp.is_transcription());
        assert!(BuiltinStep::SttDeepgram.is_transcription());
        assert!(BuiltinStep::SttOpenAi.is_transcription());
        assert!(BuiltinStep::SttGoogle.is_transcription());
        assert!(!BuiltinStep::Refine.is_transcription());
    }

    #[test]
    fn internal_step_invocation_rejects_unknown_step_name() {
        let args = vec![
            "muninn".to_string(),
            "__internal_step".to_string(),
            "unknown_step".to_string(),
        ];

        assert_eq!(maybe_handle_internal_step(&args), Some(ExitCode::FAILURE));
    }

    #[test]
    fn builtin_step_config_loader_uses_defaults_only_for_missing_config() {
        let resolved = muninn::resolve_builtin_step_config_from_load_result(
            "OpenAI provider",
            Err(muninn::ConfigError::NotFound {
                path: PathBuf::from("/tmp/missing-config.toml"),
            }),
            || "default-value".to_string(),
            |_| "resolved-value".to_string(),
        )
        .expect("missing config should fall back to defaults");

        assert_eq!(resolved, "default-value");
    }

    #[test]
    fn builtin_step_config_loader_rejects_invalid_config() {
        let resolved = muninn::resolve_builtin_step_config_from_load_result(
            "OpenAI provider",
            Err(muninn::ConfigError::Validation(
                ConfigValidationError::RefineEndpointMustNotBeEmpty,
            )),
            || "default-value".to_string(),
            |_| "resolved-value".to_string(),
        );

        assert_eq!(
            resolved,
            Err(
                "failed to load AppConfig for OpenAI provider: refine.endpoint must not be empty"
                    .to_string()
            )
        );
    }
}