oxios-kernel 1.0.2

//! Engine provider — wraps oxi-sdk's `Oxi` for the kernel.
//!
//! All provider/model resolution goes through `oxi_sdk::OxiBuilder`.
//! The `OxiosEngine` struct wraps the SDK instance and exposes a clean API
//! with support for routing, credentials, provider pooling, and multi-provider fallback.
//!
//! # Architecture
//!
//! ```text
//! OxiosEngine (OxiBuilder → Oxi)
//!   ├── resolve_model("provider/model") → Model
//!   ├── create_provider("anthropic")     → Arc<dyn Provider>
//!   ├── pooled_provider("anthropic")     → Arc<dyn Provider> (rate-limited)
//!   ├── oxi()                            → &Oxi (for AgentBuilder, etc.)
//!   └── agent(AgentConfig)               → AgentBuilder
//! ```

use anyhow::Result;
use oxi_sdk::{Oxi, OxiBuilder, ProviderPool, RateLimitPolicy};
use std::sync::Arc;

use crate::credential::CredentialStore;

/// The kernel's engine — wraps oxi-sdk's Oxi instance.
///
/// Created via [`OxiosEngine::new()`] or [`OxiosEngine::builder()`].
/// Provides access to providers, models, routing, pooling, and agent construction.
pub struct OxiosEngine {
    oxi: Oxi,
    default_model_id: String,
    /// Runtime routing control for dynamic model selection.
    routing_control: Option<oxi_sdk::RoutingControl>,
    /// Pooled providers with rate limiting.
    /// Key: provider name (e.g. "anthropic"), Value: ProviderPool wrapper.
    pools: parking_lot::RwLock<std::collections::HashMap<String, Arc<dyn oxi_sdk::Provider>>>,
}

impl OxiosEngine {
    /// Create a new engine with the given default model.
    ///
    /// Internally calls `OxiBuilder::new().with_builtins()` to load all
    /// built-in models and providers.
    pub fn new(default_model_id: impl Into<String>) -> Self {
        let model_id = default_model_id.into();
        let oxi = OxiBuilder::new().with_builtins().build();
        Self {
            oxi,
            default_model_id: model_id,
            routing_control: None,
            pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
        }
    }

    /// Create a new engine with credentials from config.
    ///
    /// Resolves API keys from CredentialStore for each known provider
    /// and injects them into the OxiBuilder. This enables the engine
    /// to create properly authenticated providers.
    ///
    /// Resolution order (per provider): env var → config.toml → ~/.oxi/auth.json
    pub fn from_config(default_model_id: impl Into<String>, config_api_key: Option<&str>) -> Self {
        let model_id = default_model_id.into();

        // Resolve the primary provider's credential
        let primary_provider = model_id
            .split_once('/')
            .map(|(p, _)| p)
            .unwrap_or("anthropic");

        let mut builder = OxiBuilder::new().with_builtins();

        // Inject credentials for all major providers via CredentialStore.
        // This ensures `create_provider()` can always build an authenticated provider.
        let providers = ["anthropic", "openai", "google", "deepseek", "xai"];
        for provider in providers {
            // Use the config-level key only for the primary provider;
            // other providers resolve from env/auth.json.
            let config_key = if provider == primary_provider {
                config_api_key
            } else {
                None
            };

            if let Some((key, source)) = CredentialStore::resolve(provider, config_key) {
                tracing::debug!(
                    provider,
                    source = ?source,
                    "Injected credential into engine"
                );
                builder = builder.api_key(provider, key);
            }
        }

        let oxi = builder.build();
        Self {
            oxi,
            default_model_id: model_id,
            routing_control: None,
            pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
        }
    }

    /// Create an engine builder for advanced configuration.
    ///
    /// Use this when you need credential injection, routing, or
    /// custom provider registration.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use oxios_kernel::engine::OxiosEngine;
    ///
    /// let engine = OxiosEngine::builder()
    ///     .default_model("anthropic/claude-sonnet-4-20250514")
    ///     .api_key("anthropic", "sk-ant-...")
    ///     .build();
    /// ```
    pub fn builder() -> OxiosEngineBuilder {
        OxiosEngineBuilder {
            inner: OxiBuilder::new().with_builtins(),
            default_model_id: "anthropic/claude-sonnet-4-20250514".to_string(),
        }
    }

    /// Get a reference to the underlying Oxi instance.
    ///
    /// Use this when you need to pass the engine to oxi-sdk APIs directly
    /// (e.g., `AgentBuilder`, `MessageBus`, `AgentGroup`).
    pub fn oxi(&self) -> &Oxi {
        &self.oxi
    }

    /// Resolve a model ID to a Model.
    pub fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
        self.oxi.resolve_model(model_id)
    }

    /// Create a provider for the given provider name.
    pub fn create_provider(&self, name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
        self.oxi.create_provider(name)
    }

    /// Get the default model ID.
    pub fn default_model_id(&self) -> &str {
        &self.default_model_id
    }

    /// Get the routing control, if routing is enabled.
    pub fn routing_control(&self) -> Option<&oxi_sdk::RoutingControl> {
        self.routing_control.as_ref()
    }

    /// Get a rate-limited provider from the pool.
    ///
    /// On first call for a provider name, creates a `ProviderPool` wrapping
    /// the base provider with the given RPM/concurrency limits.
    /// Subsequent calls return the same pooled instance.
    ///
    /// If no rate limit is needed, returns the base provider directly.
    pub fn pooled_provider(&self, name: &str, rpm: u32) -> Result<Arc<dyn oxi_sdk::Provider>> {
        // Check if already pooled.
        {
            let pools = self.pools.read();
            if let Some(pooled) = pools.get(name) {
                return Ok(pooled.clone());
            }
        }

        // Create new pool.
        let base = self.create_provider(name)?;
        let policy = RateLimitPolicy::rpm(rpm);
        let pool = ProviderPool::new(base, policy, name);
        let pooled: Arc<dyn oxi_sdk::Provider> = Arc::new(pool);

        // Cache it.
        {
            let mut pools = self.pools.write();
            pools.insert(name.to_string(), pooled.clone());
        }

        tracing::info!(provider = name, rpm, "Created provider pool");
        Ok(pooled)
    }
}

// ---------------------------------------------------------------------------
// EngineBuilder
// ---------------------------------------------------------------------------

/// Builder for creating an `OxiosEngine` with advanced configuration.
pub struct OxiosEngineBuilder {
    inner: OxiBuilder,
    default_model_id: String,
}

impl OxiosEngineBuilder {
    /// Set the default model ID.
    pub fn default_model(mut self, model_id: impl Into<String>) -> Self {
        self.default_model_id = model_id.into();
        self
    }

    /// Register an API key for a specific provider.
    pub fn api_key(self, provider: &str, key: impl Into<String>) -> Self {
        Self {
            inner: self.inner.api_key(provider, key),
            default_model_id: self.default_model_id,
        }
    }

    /// Register a full credential (API key + optional base URL).
    pub fn credential(
        self,
        provider: &str,
        api_key: impl Into<String>,
        base_url: Option<&str>,
    ) -> Self {
        Self {
            inner: self.inner.credential(provider, api_key, base_url),
            default_model_id: self.default_model_id,
        }
    }

    /// Register a custom provider.
    pub fn provider(self, name: &str, p: impl oxi_sdk::Provider + 'static) -> Self {
        Self {
            inner: self.inner.provider(name, p),
            default_model_id: self.default_model_id,
        }
    }

    /// Build the engine.
    pub fn build(self) -> OxiosEngine {
        OxiosEngine {
            oxi: self.inner.build(),
            default_model_id: self.default_model_id,
            routing_control: None,
            pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
        }
    }

    /// Build the engine with routing enabled.
    ///
    /// Returns `(OxiosEngine, RoutingControl)` for runtime routing control.
    pub fn build_with_routing(self) -> (OxiosEngine, oxi_sdk::RoutingControl) {
        use oxi_sdk::RoutingControl;

        let routing_config = oxi_sdk::routing::RoutingConfig::default();
        let routing_control = RoutingControl::new(routing_config);
        let engine = OxiosEngine {
            oxi: self.inner.build(),
            default_model_id: self.default_model_id,
            routing_control: Some(routing_control.clone()),
            pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
        };
        (engine, routing_control)
    }
}

// ---------------------------------------------------------------------------
// EngineProvider trait (for testability and dependency inversion)
// ---------------------------------------------------------------------------

/// Engine provider trait — abstracts how the kernel obtains AI providers.
///
/// Implemented by `OxiosEngine` directly. Use a mock for testing.
pub trait EngineProvider: Send + Sync {
    /// Create a provider for the given provider name.
    fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>>;

    /// Resolve a "provider/model" string to a Model.
    fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model>;

    /// Get the default model ID.
    fn default_model_id(&self) -> &str;
}

impl EngineProvider for OxiosEngine {
    fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
        self.create_provider(provider_name)
    }

    fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
        self.resolve_model(model_id)
    }

    fn default_model_id(&self) -> &str {
        &self.default_model_id
    }
}

impl std::fmt::Debug for OxiosEngine {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OxiosEngine")
            .field("default_model_id", &self.default_model_id)
            .field("routing_enabled", &self.routing_control.is_some())
            .finish()
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn test_resolve_model_with_provider_prefix() {
        let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
        let model = engine.resolve_model("openai/gpt-4o").unwrap();
        assert_eq!(model.provider, "openai");
        assert_eq!(model.id, "gpt-4o");
    }

    #[test]
    fn test_resolve_model_without_provider_prefix() {
        let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
        let model = engine.resolve_model("claude-sonnet-4-20250514").unwrap();
        assert_eq!(model.provider, "anthropic");
    }

    #[test]
    fn test_default_model_id() {
        let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
        assert_eq!(
            engine.default_model_id(),
            "anthropic/claude-sonnet-4-20250514"
        );
    }

    #[test]
    fn test_resolve_model_not_found() {
        let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
        let result = engine.resolve_model("nonexistent/model-xyz");
        assert!(result.is_err());
    }

    #[test]
    fn test_create_provider_anthropic() {
        let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
        let provider = engine.create_provider("anthropic");
        assert!(provider.is_ok());
    }

    #[test]
    fn test_create_provider_not_found() {
        let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
        let result = engine.create_provider("nonexistent_provider");
        assert!(result.is_err());
    }

    #[test]
    fn test_builder_with_credential() {
        let engine = OxiosEngine::builder()
            .default_model("openai/gpt-4o")
            .credential("openai", "sk-test", None)
            .build();
        assert_eq!(engine.default_model_id(), "openai/gpt-4o");
    }

    #[test]
    fn test_engine_provider_trait_on_engine() {
        let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
        let provider: &dyn EngineProvider = &engine;
        assert!(provider.create_provider("anthropic").is_ok());
        assert!(provider.resolve_model("openai/gpt-4o").is_ok());
    }
}