transcription.rs

use anyhow::{bail, Context, Result};
use reqwest::multipart::{Form, Part};
use reqwest::Url;

use crate::config::TranscriptionConfig;

/// Maximum upload size accepted by the Groq Whisper API (25 MB).
const MAX_AUDIO_BYTES: usize = 25 * 1024 * 1024;

/// Map file extension to MIME type for Whisper-compatible transcription APIs.
fn mime_for_audio(extension: &str) -> Option<&'static str> {
    match extension.to_ascii_lowercase().as_str() {
        "flac" => Some("audio/flac"),
        "mp3" | "mpeg" | "mpga" => Some("audio/mpeg"),
        "mp4" | "m4a" => Some("audio/mp4"),
        "ogg" | "oga" => Some("audio/ogg"),
        "opus" => Some("audio/opus"),
        "wav" => Some("audio/wav"),
        "webm" => Some("audio/webm"),
        _ => None,
    }
}

/// Normalize audio filename for Whisper-compatible APIs.
///
/// Groq validates the filename extension — `.oga` (Opus-in-Ogg) is not in
/// its accepted list, so we rewrite it to `.ogg`.
fn normalize_audio_filename(file_name: &str) -> String {
    match file_name.rsplit_once('.') {
        Some((stem, ext)) if ext.eq_ignore_ascii_case("oga") => format!("{stem}.ogg"),
        _ => file_name.to_string(),
    }
}

/// Returns `true` when `api_url` points to a Mistral endpoint.
///
/// Parses the URL and inspects the host (case-insensitive).  Falls back to
/// `false` on parse errors so the Groq default path is used.
fn is_mistral_host(api_url: &str) -> bool {
    Url::parse(api_url)
        .ok()
        .and_then(|u| u.host_str().map(|h| h.to_ascii_lowercase()))
        .map_or(false, |host| {
            host == "mistral.ai" || host.ends_with(".mistral.ai")
        })
}

/// Transcribe audio bytes via a Whisper-compatible transcription API.
///
/// Supports Groq Whisper (default) and Mistral Voxtral endpoints.
/// The provider is detected from `config.api_url` by inspecting the host.
///
/// API key resolution order:
/// 1. Explicit `config.api_key` (highest priority).
/// 2. `MISTRAL_API_KEY` env var (when the endpoint is `*.mistral.ai`).
/// 3. `GROQ_API_KEY` env var (all other endpoints).
///
/// The caller is responsible for enforcing duration limits *before*
/// downloading the file; this function enforces the byte-size cap.
pub async fn transcribe_audio(
    audio_data: Vec<u8>,
    file_name: &str,
    config: &TranscriptionConfig,
) -> Result<String> {
    if audio_data.len() > MAX_AUDIO_BYTES {
        bail!(
            "Audio file too large ({} bytes, max {MAX_AUDIO_BYTES})",
            audio_data.len()
        );
    }

    let normalized_name = normalize_audio_filename(file_name);
    let extension = normalized_name
        .rsplit_once('.')
        .map(|(_, e)| e)
        .unwrap_or("");
    let mime = mime_for_audio(extension).ok_or_else(|| {
        anyhow::anyhow!(
            "Unsupported audio format '.{extension}' — accepted: flac, mp3, mp4, mpeg, mpga, m4a, ogg, opus, wav, webm"
        )
    })?;

    let mistral = is_mistral_host(&config.api_url);

    let api_key = config
        .api_key
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToOwned::to_owned)
        .or_else(|| {
            let var = if mistral {
                "MISTRAL_API_KEY"
            } else {
                "GROQ_API_KEY"
            };
            std::env::var(var)
                .ok()
                .map(|value| value.trim().to_string())
                .filter(|value| !value.is_empty())
        })
        .context(
            "Missing transcription API key: set [transcription].api_key, MISTRAL_API_KEY, or GROQ_API_KEY environment variable",
        )?;

    let proxy_name = if mistral {
        "transcription.mistral"
    } else {
        "transcription.groq"
    };
    let client = crate::config::build_runtime_proxy_client(proxy_name);

    let file_part = Part::bytes(audio_data)
        .file_name(normalized_name)
        .mime_str(mime)?;

    let mut form = Form::new()
        .part("file", file_part)
        .text("model", config.model.clone())
        .text("response_format", "json");

    if let Some(ref lang) = config.language {
        form = form.text("language", lang.clone());
    }

    let resp = client
        .post(&config.api_url)
        .bearer_auth(&api_key)
        .multipart(form)
        .send()
        .await
        .context("Failed to send transcription request")?;

    let status = resp.status();
    let body: serde_json::Value = resp
        .json()
        .await
        .context("Failed to parse transcription response")?;

    if !status.is_success() {
        let error_msg = body["error"]["message"].as_str().unwrap_or("unknown error");
        bail!("Transcription API error ({}): {}", status, error_msg);
    }

    let text = body["text"]
        .as_str()
        .context("Transcription response missing 'text' field")?
        .to_string();

    Ok(text)
}

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

    #[tokio::test]
    async fn rejects_oversized_audio() {
        let big = vec![0u8; MAX_AUDIO_BYTES + 1];
        let config = TranscriptionConfig::default();

        let err = transcribe_audio(big, "test.ogg", &config)
            .await
            .unwrap_err();
        assert!(
            err.to_string().contains("too large"),
            "expected size error, got: {err}"
        );
    }

    #[tokio::test]
    async fn rejects_missing_api_key() {
        // Ensure the key is absent for this test
        std::env::remove_var("GROQ_API_KEY");

        let data = vec![0u8; 100];
        let config = TranscriptionConfig::default();

        let err = transcribe_audio(data, "test.ogg", &config)
            .await
            .unwrap_err();
        assert!(
            err.to_string().contains("Missing transcription API key"),
            "expected missing-key error, got: {err}"
        );
    }

    #[tokio::test]
    async fn uses_config_api_key_without_groq_env() {
        std::env::remove_var("GROQ_API_KEY");
        std::env::remove_var("MISTRAL_API_KEY");

        let mut config = TranscriptionConfig::default();
        config.api_key = Some("explicit-key".to_string());

        // Will fail on the HTTP request (no real server), but should NOT
        // fail on API key resolution.
        let err = transcribe_audio(vec![0u8; 100], "test.ogg", &config)
            .await
            .unwrap_err();
        assert!(
            !err.to_string().contains("Missing transcription API key"),
            "should not fail on key resolution when config.api_key is set, got: {err}"
        );
    }

    #[tokio::test]
    async fn mistral_url_falls_back_to_mistral_env_key() {
        std::env::remove_var("GROQ_API_KEY");
        std::env::remove_var("MISTRAL_API_KEY");

        let mut config = TranscriptionConfig::default();
        config.api_url = "https://api.mistral.ai/v1/audio/transcriptions".to_string();
        config.api_key = None;

        // Without MISTRAL_API_KEY set, should get the missing-key error.
        let err = transcribe_audio(vec![0u8; 100], "test.ogg", &config)
            .await
            .unwrap_err();
        assert!(
            err.to_string().contains("Missing transcription API key"),
            "expected missing-key error for Mistral URL without env key, got: {err}"
        );
    }

    #[tokio::test]
    async fn whitespace_only_api_key_is_rejected() {
        std::env::remove_var("GROQ_API_KEY");
        std::env::remove_var("MISTRAL_API_KEY");

        let mut config = TranscriptionConfig::default();
        config.api_key = Some("   ".to_string());

        let err = transcribe_audio(vec![0u8; 100], "test.ogg", &config)
            .await
            .unwrap_err();
        assert!(
            err.to_string().contains("Missing transcription API key"),
            "whitespace-only api_key should be treated as missing, got: {err}"
        );
    }

    // ── is_mistral_host tests ───────────────────────────────────────

    #[test]
    fn is_mistral_host_detects_api_subdomain() {
        assert!(is_mistral_host(
            "https://api.mistral.ai/v1/audio/transcriptions"
        ));
    }

    #[test]
    fn is_mistral_host_detects_bare_domain() {
        assert!(is_mistral_host("https://mistral.ai/endpoint"));
    }

    #[test]
    fn is_mistral_host_case_insensitive() {
        assert!(is_mistral_host(
            "https://API.MISTRAL.AI/v1/audio/transcriptions"
        ));
    }

    #[test]
    fn is_mistral_host_rejects_groq_url() {
        assert!(!is_mistral_host(
            "https://api.groq.com/openai/v1/audio/transcriptions"
        ));
    }

    #[test]
    fn is_mistral_host_rejects_spoofed_path() {
        // "mistral.ai" in path but not in host
        assert!(!is_mistral_host(
            "https://evil.com/mistral.ai/v1/audio/transcriptions"
        ));
    }

    #[test]
    fn is_mistral_host_returns_false_for_invalid_url() {
        assert!(!is_mistral_host("not-a-url"));
        assert!(!is_mistral_host(""));
    }

    // ── MIME / filename tests ───────────────────────────────────────

    #[test]
    fn mime_for_audio_maps_accepted_formats() {
        let cases = [
            ("flac", "audio/flac"),
            ("mp3", "audio/mpeg"),
            ("mpeg", "audio/mpeg"),
            ("mpga", "audio/mpeg"),
            ("mp4", "audio/mp4"),
            ("m4a", "audio/mp4"),
            ("ogg", "audio/ogg"),
            ("oga", "audio/ogg"),
            ("opus", "audio/opus"),
            ("wav", "audio/wav"),
            ("webm", "audio/webm"),
        ];
        for (ext, expected) in cases {
            assert_eq!(
                mime_for_audio(ext),
                Some(expected),
                "failed for extension: {ext}"
            );
        }
    }

    #[test]
    fn mime_for_audio_case_insensitive() {
        assert_eq!(mime_for_audio("OGG"), Some("audio/ogg"));
        assert_eq!(mime_for_audio("MP3"), Some("audio/mpeg"));
        assert_eq!(mime_for_audio("Opus"), Some("audio/opus"));
    }

    #[test]
    fn mime_for_audio_rejects_unknown() {
        assert_eq!(mime_for_audio("txt"), None);
        assert_eq!(mime_for_audio("pdf"), None);
        assert_eq!(mime_for_audio("aac"), None);
        assert_eq!(mime_for_audio(""), None);
    }

    #[test]
    fn normalize_audio_filename_rewrites_oga() {
        assert_eq!(normalize_audio_filename("voice.oga"), "voice.ogg");
        assert_eq!(normalize_audio_filename("file.OGA"), "file.ogg");
    }

    #[test]
    fn normalize_audio_filename_preserves_accepted() {
        assert_eq!(normalize_audio_filename("voice.ogg"), "voice.ogg");
        assert_eq!(normalize_audio_filename("track.mp3"), "track.mp3");
        assert_eq!(normalize_audio_filename("clip.opus"), "clip.opus");
    }

    #[test]
    fn normalize_audio_filename_no_extension() {
        assert_eq!(normalize_audio_filename("voice"), "voice");
    }

    #[tokio::test]
    async fn rejects_unsupported_audio_format() {
        let data = vec![0u8; 100];
        let config = TranscriptionConfig::default();

        let err = transcribe_audio(data, "recording.aac", &config)
            .await
            .unwrap_err();
        let msg = err.to_string();
        assert!(
            msg.contains("Unsupported audio format"),
            "expected unsupported-format error, got: {msg}"
        );
        assert!(
            msg.contains(".aac"),
            "error should mention the rejected extension, got: {msg}"
        );
    }
}