herolib-ai 0.3.13

//! Rhai scripting bindings for herolib-ai.
//!
//! This module provides AI functions for Rhai scripts including chat, embeddings, and transcription.
//!
//! ## Chat Functions
//!
//! - `ai_chat(prompt)` - Simple chat with default model
//! - `ai_chat_with_model(model, prompt)` - Chat with specific model
//! - `ai_chat_with_system(model, system, prompt)` - Chat with system message
//!
//! ## Embedding Functions
//!
//! - `ai_embed(text)` - Generate embedding for text
//! - `ai_embed_with_model(model, text)` - Generate embedding with specific model
//! - `ai_embed_batch(texts)` - Generate embeddings for multiple texts
//!
//! ## Transcription Functions
//!
//! - `ai_transcribe(file_path)` - Transcribe audio file
//! - `ai_transcribe_with_model(model, file_path)` - Transcribe with specific model
//! - `ai_transcribe_with_options(model, file_path, language)` - Transcribe with options
//!
//! ## Model Constants
//!
//! Chat models: `"llama3_3_70b"`, `"llama3_1_70b"`, `"qwen2_5_coder_32b"`, `"deepseek_v3"`, `"gpt_oss_120b"`, etc.
//! Embedding models: `"text_embedding_3_small"`, `"qwen3_embedding_8b"`
//! Transcription models: `"whisper_large_v3_turbo"`, `"whisper_large_v3"`
//!
//! ## Example
//!
//! ```rhai
//! // Simple chat
//! let response = ai_chat("What is 2 + 2?");
//! print(response);
//!
//! // Chat with specific model and system prompt
//! let response = ai_chat_with_system("llama3_3_70b", "You are a helpful assistant", "Hello!");
//! print(response);
//!
//! // Generate embedding
//! let embedding = ai_embed("Hello, world!");
//! print(`Embedding dimensions: ${embedding.len()}`);
//!
//! // Transcribe audio
//! let text = ai_transcribe("/path/to/audio.mp3");
//! print(text);
//! ```

use rhai::{Array, Engine, EvalAltResult};
use std::path::Path;

use crate::client::AiClient;
use crate::embedding::EmbeddingModel;
use crate::model::Model;
use crate::transcription::{TranscriptionModel, TranscriptionOptions};
use crate::types::Message;

/// Helper to convert errors to Rhai errors.
fn to_rhai_error<E: std::fmt::Display>(msg: &str, e: E) -> Box<EvalAltResult> {
    Box::new(EvalAltResult::ErrorRuntime(
        format!("{}: {}", msg, e).into(),
        rhai::Position::NONE,
    ))
}

/// Parse model string to Model enum.
fn parse_model(model: &str) -> Result<Model, Box<EvalAltResult>> {
    match model.to_lowercase().replace("-", "_").as_str() {
        "llama3_3_70b" | "llama3.3_70b" => Ok(Model::Llama3_3_70B),
        "llama3_1_70b" | "llama3.1_70b" => Ok(Model::Llama3_1_70B),
        "llama3_1_8b" | "llama3.1_8b" => Ok(Model::Llama3_1_8B),
        "qwen2_5_coder_32b" | "qwen2.5_coder_32b" => Ok(Model::Qwen2_5Coder32B),
        "deepseek_coder_v2_5" | "deepseek_coder" => Ok(Model::DeepSeekCoderV2_5),
        "deepseek_v3" | "deepseek" => Ok(Model::DeepSeekV3),
        "llama3_1_405b" | "llama3.1_405b" => Ok(Model::Llama3_1_405B),
        "mixtral_8x7b" | "mixtral" => Ok(Model::Mixtral8x7B),
        "llama3_2_90b_vision" | "llama3.2_90b_vision" => Ok(Model::Llama3_2_90BVision),
        "llama3_2_11b_vision" | "llama3.2_11b_vision" => Ok(Model::Llama3_2_11BVision),
        "nemotron_nano_30b" | "nemotron" => Ok(Model::NemotronNano30B),
        "gpt_oss_120b" | "gpt_oss" | "gptoss" => Ok(Model::GptOss120B),
        _ => Err(Box::new(EvalAltResult::ErrorRuntime(
            format!("Unknown model: {}. Available: llama3_3_70b, llama3_1_70b, llama3_1_8b, qwen2_5_coder_32b, deepseek_v3, mixtral_8x7b, nemotron_nano_30b, gpt_oss_120b", model).into(),
            rhai::Position::NONE,
        ))),
    }
}

/// Parse embedding model string to EmbeddingModel enum.
fn parse_embedding_model(model: &str) -> Result<EmbeddingModel, Box<EvalAltResult>> {
    match model.to_lowercase().replace("-", "_").as_str() {
        "text_embedding_3_small" | "openai_small" => Ok(EmbeddingModel::TextEmbedding3Small),
        "qwen3_embedding_8b" | "qwen_embedding" => Ok(EmbeddingModel::Qwen3Embedding8B),
        _ => Err(Box::new(EvalAltResult::ErrorRuntime(
            format!(
                "Unknown embedding model: {}. Available: text_embedding_3_small, qwen3_embedding_8b",
                model
            )
            .into(),
            rhai::Position::NONE,
        ))),
    }
}

/// Parse transcription model string to TranscriptionModel enum.
fn parse_transcription_model(model: &str) -> Result<TranscriptionModel, Box<EvalAltResult>> {
    match model.to_lowercase().replace("-", "_").as_str() {
        "whisper_large_v3_turbo" | "whisper_turbo" => Ok(TranscriptionModel::WhisperLargeV3Turbo),
        "whisper_large_v3" | "whisper" => Ok(TranscriptionModel::WhisperLargeV3),
        _ => Err(Box::new(EvalAltResult::ErrorRuntime(
            format!(
                "Unknown transcription model: {}. Available: whisper_large_v3_turbo, whisper_large_v3",
                model
            )
            .into(),
            rhai::Position::NONE,
        ))),
    }
}

/// Get client from environment, returning error if no providers configured.
fn get_client() -> Result<AiClient, Box<EvalAltResult>> {
    let client = AiClient::from_env();
    if !client.has_providers() {
        return Err(Box::new(EvalAltResult::ErrorRuntime(
            "No AI providers configured. Set GROQ_API_KEY, OPENROUTER_API_KEY, or SAMBANOVA_API_KEY environment variable.".into(),
            rhai::Position::NONE,
        )));
    }
    Ok(client)
}

// ============================================================================
// Chat Functions
// ============================================================================

/// Simple chat with default model (Llama 3.3 70B).
fn rhai_ai_chat(prompt: &str) -> Result<String, Box<EvalAltResult>> {
    let client = get_client()?;
    let messages = vec![Message::user(prompt)];

    let response = client
        .chat(Model::default_general(), messages)
        .map_err(|e| to_rhai_error("Chat failed", e))?;

    response
        .content()
        .map(|s| s.to_string())
        .ok_or_else(|| to_rhai_error("No content in response", "empty response"))
}

/// Chat with specific model.
fn rhai_ai_chat_with_model(model: &str, prompt: &str) -> Result<String, Box<EvalAltResult>> {
    let client = get_client()?;
    let model = parse_model(model)?;
    let messages = vec![Message::user(prompt)];

    let response = client
        .chat(model, messages)
        .map_err(|e| to_rhai_error("Chat failed", e))?;

    response
        .content()
        .map(|s| s.to_string())
        .ok_or_else(|| to_rhai_error("No content in response", "empty response"))
}

/// Chat with system message.
fn rhai_ai_chat_with_system(
    model: &str,
    system: &str,
    prompt: &str,
) -> Result<String, Box<EvalAltResult>> {
    let client = get_client()?;
    let model = parse_model(model)?;
    let messages = vec![Message::system(system), Message::user(prompt)];

    let response = client
        .chat(model, messages)
        .map_err(|e| to_rhai_error("Chat failed", e))?;

    response
        .content()
        .map(|s| s.to_string())
        .ok_or_else(|| to_rhai_error("No content in response", "empty response"))
}

// ============================================================================
// Embedding Functions
// ============================================================================

/// Generate embedding for text with default model.
fn rhai_ai_embed(text: &str) -> Result<Array, Box<EvalAltResult>> {
    let client = get_client()?;

    let response = client
        .embed(EmbeddingModel::default(), text)
        .map_err(|e| to_rhai_error("Embedding failed", e))?;

    let embedding = response
        .embedding()
        .ok_or_else(|| to_rhai_error("No embedding in response", "empty response"))?;

    Ok(embedding
        .iter()
        .map(|&f| rhai::Dynamic::from(f as f64))
        .collect())
}

/// Generate embedding with specific model.
fn rhai_ai_embed_with_model(model: &str, text: &str) -> Result<Array, Box<EvalAltResult>> {
    let client = get_client()?;
    let model = parse_embedding_model(model)?;

    let response = client
        .embed(model, text)
        .map_err(|e| to_rhai_error("Embedding failed", e))?;

    let embedding = response
        .embedding()
        .ok_or_else(|| to_rhai_error("No embedding in response", "empty response"))?;

    Ok(embedding
        .iter()
        .map(|&f| rhai::Dynamic::from(f as f64))
        .collect())
}

/// Generate embeddings for multiple texts.
fn rhai_ai_embed_batch(texts: Array) -> Result<Array, Box<EvalAltResult>> {
    let client = get_client()?;

    let texts: Vec<String> = texts
        .into_iter()
        .map(|v| {
            v.into_string()
                .map_err(|_| to_rhai_error("Invalid text", "expected string"))
        })
        .collect::<Result<Vec<_>, _>>()?;

    let response = client
        .embed_batch(EmbeddingModel::default(), texts)
        .map_err(|e| to_rhai_error("Batch embedding failed", e))?;

    let embeddings: Array = response
        .embeddings()
        .iter()
        .map(|emb| {
            rhai::Dynamic::from(
                emb.iter()
                    .map(|&f| rhai::Dynamic::from(f as f64))
                    .collect::<Array>(),
            )
        })
        .collect();

    Ok(embeddings)
}

// ============================================================================
// Transcription Functions
// ============================================================================

/// Transcribe audio file with default model.
fn rhai_ai_transcribe(file_path: &str) -> Result<String, Box<EvalAltResult>> {
    let client = get_client()?;
    let path = Path::new(file_path);

    if !path.exists() {
        return Err(Box::new(EvalAltResult::ErrorRuntime(
            format!("Audio file not found: {}", file_path).into(),
            rhai::Position::NONE,
        )));
    }

    let response = client
        .transcribe_file(TranscriptionModel::default(), path)
        .map_err(|e| to_rhai_error("Transcription failed", e))?;

    Ok(response.text)
}

/// Transcribe audio file with specific model.
fn rhai_ai_transcribe_with_model(
    model: &str,
    file_path: &str,
) -> Result<String, Box<EvalAltResult>> {
    let client = get_client()?;
    let model = parse_transcription_model(model)?;
    let path = Path::new(file_path);

    if !path.exists() {
        return Err(Box::new(EvalAltResult::ErrorRuntime(
            format!("Audio file not found: {}", file_path).into(),
            rhai::Position::NONE,
        )));
    }

    let response = client
        .transcribe_file(model, path)
        .map_err(|e| to_rhai_error("Transcription failed", e))?;

    Ok(response.text)
}

/// Transcribe audio file with options.
fn rhai_ai_transcribe_with_options(
    model: &str,
    file_path: &str,
    language: &str,
) -> Result<String, Box<EvalAltResult>> {
    let client = get_client()?;
    let model = parse_transcription_model(model)?;
    let path = Path::new(file_path);

    if !path.exists() {
        return Err(Box::new(EvalAltResult::ErrorRuntime(
            format!("Audio file not found: {}", file_path).into(),
            rhai::Position::NONE,
        )));
    }

    let options = TranscriptionOptions::new().with_language(language);

    let response = client
        .transcribe_file_with_options(model, path, options)
        .map_err(|e| to_rhai_error("Transcription failed", e))?;

    Ok(response.text)
}

/// List available chat models.
fn rhai_ai_models() -> Array {
    Model::all()
        .iter()
        .map(|m| rhai::Dynamic::from(m.name().to_string()))
        .collect()
}

/// List available embedding models.
fn rhai_ai_embedding_models() -> Array {
    EmbeddingModel::all()
        .iter()
        .map(|m| rhai::Dynamic::from(m.name().to_string()))
        .collect()
}

/// List available transcription models.
fn rhai_ai_transcription_models() -> Array {
    TranscriptionModel::all()
        .iter()
        .map(|m| rhai::Dynamic::from(m.name().to_string()))
        .collect()
}

// ============================================================================
// Registration
// ============================================================================

/// Register all AI functions with a Rhai engine.
///
/// ## Functions Registered
///
/// ### Chat
/// - `ai_chat(prompt)` - Simple chat with default model
/// - `ai_chat_with_model(model, prompt)` - Chat with specific model
/// - `ai_chat_with_system(model, system, prompt)` - Chat with system message
///
/// ### Embeddings
/// - `ai_embed(text)` - Generate embedding for text
/// - `ai_embed_with_model(model, text)` - Generate embedding with specific model
/// - `ai_embed_batch(texts)` - Generate embeddings for multiple texts
///
/// ### Transcription
/// - `ai_transcribe(file_path)` - Transcribe audio file
/// - `ai_transcribe_with_model(model, file_path)` - Transcribe with specific model
/// - `ai_transcribe_with_options(model, file_path, language)` - Transcribe with options
///
/// ### Utilities
/// - `ai_models()` - List available chat models
/// - `ai_embedding_models()` - List available embedding models
/// - `ai_transcription_models()` - List available transcription models
///
/// ## Example
///
/// ```rust,ignore
/// use rhai::Engine;
/// use herolib_ai::rhai::register;
///
/// let mut engine = Engine::new();
/// register(&mut engine).unwrap();
/// ```
pub fn register(engine: &mut Engine) -> Result<(), Box<EvalAltResult>> {
    // Chat functions
    engine.register_fn("ai_chat", rhai_ai_chat);
    engine.register_fn("ai_chat_with_model", rhai_ai_chat_with_model);
    engine.register_fn("ai_chat_with_system", rhai_ai_chat_with_system);

    // Embedding functions
    engine.register_fn("ai_embed", rhai_ai_embed);
    engine.register_fn("ai_embed_with_model", rhai_ai_embed_with_model);
    engine.register_fn("ai_embed_batch", rhai_ai_embed_batch);

    // Transcription functions
    engine.register_fn("ai_transcribe", rhai_ai_transcribe);
    engine.register_fn("ai_transcribe_with_model", rhai_ai_transcribe_with_model);
    engine.register_fn(
        "ai_transcribe_with_options",
        rhai_ai_transcribe_with_options,
    );

    // Utility functions
    engine.register_fn("ai_models", rhai_ai_models);
    engine.register_fn("ai_embedding_models", rhai_ai_embedding_models);
    engine.register_fn("ai_transcription_models", rhai_ai_transcription_models);

    Ok(())
}

/// Alias for register() for consistency with other modules.
pub fn register_ai_module(engine: &mut Engine) -> Result<(), Box<EvalAltResult>> {
    register(engine)
}

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

    #[test]
    fn test_parse_model() {
        assert!(parse_model("llama3_3_70b").is_ok());
        assert!(parse_model("deepseek_v3").is_ok());
        assert!(parse_model("invalid_model").is_err());
    }

    #[test]
    fn test_parse_embedding_model() {
        assert!(parse_embedding_model("text_embedding_3_small").is_ok());
        assert!(parse_embedding_model("qwen3_embedding_8b").is_ok());
        assert!(parse_embedding_model("invalid").is_err());
    }

    #[test]
    fn test_parse_transcription_model() {
        assert!(parse_transcription_model("whisper_large_v3_turbo").is_ok());
        assert!(parse_transcription_model("whisper_large_v3").is_ok());
        assert!(parse_transcription_model("invalid").is_err());
    }

    #[test]
    fn test_register() {
        let mut engine = Engine::new();
        assert!(register(&mut engine).is_ok());
    }

    #[test]
    fn test_ai_models_function() {
        let mut engine = Engine::new();
        register(&mut engine).unwrap();

        let result = engine.eval::<Array>("ai_models()").unwrap();
        assert!(!result.is_empty());
    }

    #[test]
    fn test_ai_embedding_models_function() {
        let mut engine = Engine::new();
        register(&mut engine).unwrap();

        let result = engine.eval::<Array>("ai_embedding_models()").unwrap();
        assert!(!result.is_empty());
    }

    #[test]
    fn test_ai_transcription_models_function() {
        let mut engine = Engine::new();
        register(&mut engine).unwrap();

        let result = engine.eval::<Array>("ai_transcription_models()").unwrap();
        assert!(!result.is_empty());
    }
}