bitrouter 0.24.0

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;

use bitrouter_config::{ApiProtocol, ProviderConfig};
use bitrouter_core::{
    errors::{BitrouterError, Result},
    models::language::language_model::DynLanguageModel,
    routers::{router::LanguageModelRouter, routing_table::RoutingTarget},
    tools::provider::DynToolProvider,
};
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use reqwest_middleware::ClientWithMiddleware;

use super::copilot::{
    COPILOT_PROVIDER, anthropic_api_base, copilot_default_headers, is_anthropic_model,
};

#[cfg(feature = "mcp")]
use bitrouter_providers::mcp::client::upstream::UpstreamConnection;

/// A model router backed by `reqwest` that instantiates concrete provider
/// model objects on demand from [`ProviderConfig`] entries.
pub struct Router {
    client: ClientWithMiddleware,
    providers: HashMap<String, ProviderConfig>,
    /// Path to the OAuth token store file (`tokens.json`).
    /// When set, OAuth-authenticated providers resolve their API key from
    /// the token store at request time.
    token_store_path: Option<PathBuf>,
}

impl Router {
    pub fn new(client: ClientWithMiddleware, providers: HashMap<String, ProviderConfig>) -> Self {
        Self {
            client,
            providers,
            token_store_path: None,
        }
    }

    /// Set the path to the OAuth token store.
    pub fn with_token_store(mut self, path: PathBuf) -> Self {
        self.token_store_path = Some(path);
        self
    }

    /// Resolve the API key for a provider, checking the OAuth token store
    /// when the provider uses `auth.type: oauth`.
    fn resolve_api_key(&self, provider_name: &str, provider: &ProviderConfig) -> String {
        use bitrouter_config::AuthConfig;
        if matches!(provider.auth, Some(AuthConfig::OAuth { .. }))
            && let Some(ref path) = self.token_store_path
        {
            let store = crate::auth::token_store::TokenStore::load(path);
            if let Some(token) = store.get(provider_name) {
                return token.access_token.clone();
            }
        }
        provider.api_key.clone().unwrap_or_default()
    }

    fn build_openai_config(
        &self,
        provider_name: &str,
        provider: &ProviderConfig,
    ) -> Result<OpenAiConfig> {
        let api_key = self.resolve_api_key(provider_name, provider);
        let base_url = provider
            .api_base
            .clone()
            .unwrap_or_else(|| "https://api.openai.com/v1".into());

        let default_headers = parse_headers(provider.default_headers.as_ref())?;

        Ok(OpenAiConfig {
            api_key,
            base_url,
            organization: None,
            project: None,
            default_headers,
        })
    }

    fn build_anthropic_config(
        &self,
        provider_name: &str,
        provider: &ProviderConfig,
    ) -> Result<AnthropicConfig> {
        let api_key = self.resolve_api_key(provider_name, provider);
        let base_url = provider
            .api_base
            .clone()
            .unwrap_or_else(|| "https://api.anthropic.com".into());

        let default_headers = parse_headers(provider.default_headers.as_ref())?;

        Ok(AnthropicConfig {
            api_key,
            base_url,
            api_version: "2023-06-01".into(),
            default_headers,
        })
    }

    fn build_google_config(&self, provider: &ProviderConfig) -> Result<GoogleConfig> {
        let api_key = provider.api_key.clone().unwrap_or_default();
        let base_url = provider
            .api_base
            .clone()
            .unwrap_or_else(|| "https://generativelanguage.googleapis.com".into());

        let default_headers = parse_headers(provider.default_headers.as_ref())?;

        Ok(GoogleConfig {
            api_key,
            base_url,
            default_headers,
        })
    }
}

impl LanguageModelRouter for Router {
    async fn route_model(&self, target: RoutingTarget) -> Result<Box<DynLanguageModel<'static>>> {
        let provider = self.providers.get(&target.provider_name).ok_or_else(|| {
            BitrouterError::invalid_request(
                None,
                format!("unknown provider: {}", target.provider_name),
                None,
            )
        })?;

        let is_copilot = target.provider_name == COPILOT_PROVIDER;

        // Phase 3: Copilot Claude models use the Anthropic Messages API.
        if is_copilot && is_anthropic_model(&target.service_id) {
            let api_key = self.resolve_api_key(&target.provider_name, provider);
            let base_url = anthropic_api_base(provider.api_base.as_deref());
            let mut default_headers = parse_headers(provider.default_headers.as_ref())?;
            merge_copilot_headers(&mut default_headers)?;

            let config = AnthropicConfig {
                api_key,
                base_url,
                api_version: "2023-06-01".into(),
                default_headers,
            };
            let model =
                AnthropicMessagesModel::with_client(target.service_id, self.client.clone(), config);
            return Ok(DynLanguageModel::new_box(model));
        }

        // Phase 2: Inject Copilot-specific headers for non-Claude models.
        let copilot_headers = if is_copilot {
            Some(copilot_default_headers())
        } else {
            None
        };

        match target.api_protocol {
            ApiProtocol::Openai => {
                let mut config = self.build_openai_config(&target.provider_name, provider)?;
                if let Some(ref extra) = copilot_headers {
                    merge_into_header_map(&mut config.default_headers, extra)?;
                }
                let model = OpenAiChatCompletionsModel::with_client(
                    target.service_id,
                    self.client.clone(),
                    config,
                );
                Ok(DynLanguageModel::new_box(model))
            }
            ApiProtocol::Anthropic => {
                let config = self.build_anthropic_config(&target.provider_name, provider)?;
                let model = AnthropicMessagesModel::with_client(
                    target.service_id,
                    self.client.clone(),
                    config,
                );
                Ok(DynLanguageModel::new_box(model))
            }
            ApiProtocol::Google => {
                let config = self.build_google_config(provider)?;
                let model = GoogleGenerativeAiModel::with_client(
                    target.service_id,
                    self.client.clone(),
                    config,
                );
                Ok(DynLanguageModel::new_box(model))
            }
            ApiProtocol::Mcp | ApiProtocol::Rest | ApiProtocol::Acp => {
                Err(BitrouterError::invalid_request(
                    Some(&target.provider_name),
                    format!(
                        "provider '{}' uses protocol '{}' which cannot serve models",
                        target.provider_name, target.api_protocol
                    ),
                    None,
                ))
            }
        }
    }
}

/// Merge Copilot default headers into an existing [`HeaderMap`].
fn merge_copilot_headers(map: &mut HeaderMap) -> Result<()> {
    merge_into_header_map(map, &copilot_default_headers())
}

/// Merge string-keyed headers into a [`HeaderMap`].
fn merge_into_header_map(map: &mut HeaderMap, extra: &HashMap<String, String>) -> Result<()> {
    for (k, v) in extra {
        let name = HeaderName::from_bytes(k.as_bytes()).map_err(|e| {
            BitrouterError::invalid_request(None, format!("invalid header name '{k}': {e}"), None)
        })?;
        let value = HeaderValue::from_str(v).map_err(|e| {
            BitrouterError::invalid_request(
                None,
                format!("invalid header value for '{k}': {e}"),
                None,
            )
        })?;
        map.insert(name, value);
    }
    Ok(())
}

fn parse_headers(headers: Option<&HashMap<String, String>>) -> Result<HeaderMap> {
    let mut map = HeaderMap::new();
    if let Some(h) = headers {
        for (k, v) in h {
            let name = HeaderName::from_bytes(k.as_bytes()).map_err(|e| {
                BitrouterError::invalid_request(
                    None,
                    format!("invalid header name '{k}': {e}"),
                    None,
                )
            })?;
            let value = HeaderValue::from_str(v).map_err(|e| {
                BitrouterError::invalid_request(
                    None,
                    format!("invalid header value for '{k}': {e}"),
                    None,
                )
            })?;
            map.insert(name, value);
        }
    }
    Ok(map)
}

// Re-export provider types under short aliases for readability.
use bitrouter_providers::anthropic::messages::provider::{AnthropicConfig, AnthropicMessagesModel};
use bitrouter_providers::google::generate_content::provider::{
    GoogleConfig, GoogleGenerativeAiModel,
};
use bitrouter_providers::openai::chat::provider::{OpenAiChatCompletionsModel, OpenAiConfig};

// ── Tool router ──────────────────────────────────────────────────

/// A lazy tool router that constructs providers on-demand from config,
/// mirroring how [`Router`] constructs language models on-demand.
///
/// REST providers are stateless and constructed per-call.
/// MCP providers are looked up from a pre-built connection pool.
pub struct LazyToolRouter {
    providers: HashMap<String, ProviderConfig>,
    #[cfg(feature = "mcp")]
    mcp_pool: Arc<HashMap<String, Arc<UpstreamConnection>>>,
    #[cfg(feature = "rest")]
    client: Arc<reqwest::Client>,
}

impl LazyToolRouter {
    pub fn new(
        providers: HashMap<String, ProviderConfig>,
        #[cfg(feature = "mcp")] mcp_pool: Arc<HashMap<String, Arc<UpstreamConnection>>>,
        #[cfg(feature = "rest")] client: Arc<reqwest::Client>,
        #[cfg(not(feature = "rest"))] _client: Arc<reqwest::Client>,
    ) -> Self {
        Self {
            providers,
            #[cfg(feature = "mcp")]
            mcp_pool,
            #[cfg(feature = "rest")]
            client,
        }
    }

    /// Returns `true` if there are any tool-capable providers.
    pub fn has_providers(&self) -> bool {
        #[cfg(feature = "mcp")]
        let has_mcp = !self.mcp_pool.is_empty();
        #[cfg(not(feature = "mcp"))]
        let has_mcp = false;

        has_mcp
            || self
                .providers
                .values()
                .any(|p| p.api_protocol == Some(ApiProtocol::Rest))
    }
}

impl bitrouter_core::routers::router::ToolRouter for LazyToolRouter {
    async fn route_tool(&self, target: RoutingTarget) -> Result<Box<DynToolProvider<'static>>> {
        match target.api_protocol {
            #[cfg(feature = "rest")]
            ApiProtocol::Rest => {
                let provider = self.providers.get(&target.provider_name).ok_or_else(|| {
                    BitrouterError::invalid_request(
                        None,
                        format!("unknown REST provider: {}", target.provider_name),
                        None,
                    )
                })?;
                let api_base = provider.api_base.clone().ok_or_else(|| {
                    BitrouterError::invalid_request(
                        None,
                        format!("REST provider '{}' has no api_base", target.provider_name),
                        None,
                    )
                })?;
                let auth_header = resolve_auth_header(provider);
                let p = bitrouter_providers::rest::provider::RestToolProvider::new(
                    target.provider_name,
                    api_base,
                    auth_header,
                    self.client.clone(),
                );
                Ok(DynToolProvider::new_box(p))
            }
            #[cfg(not(feature = "rest"))]
            ApiProtocol::Rest => Err(BitrouterError::invalid_request(
                None,
                format!(
                    "REST protocol not available (feature disabled) for provider '{}'",
                    target.provider_name
                ),
                None,
            )),
            #[cfg(feature = "mcp")]
            ApiProtocol::Mcp => {
                let conn = self.mcp_pool.get(&target.provider_name).ok_or_else(|| {
                    BitrouterError::invalid_request(
                        None,
                        format!("no MCP connection for provider '{}'", target.provider_name),
                        None,
                    )
                })?;
                Ok(DynToolProvider::new_box(McpToolProviderAdapter(
                    Arc::clone(conn),
                )))
            }
            #[cfg(not(feature = "mcp"))]
            ApiProtocol::Mcp => Err(BitrouterError::invalid_request(
                None,
                format!(
                    "MCP protocol not available (feature disabled) for provider '{}'",
                    target.provider_name
                ),
                None,
            )),
            other => Err(BitrouterError::invalid_request(
                None,
                format!(
                    "protocol '{}' cannot serve tools for provider '{}'",
                    other, target.provider_name
                ),
                None,
            )),
        }
    }
}

// ── MCP → ToolProvider adapter ───────────────────────────────────

/// Thin wrapper that delegates [`ToolProvider`] to an `Arc<UpstreamConnection>`.
#[cfg(feature = "mcp")]
struct McpToolProviderAdapter(Arc<UpstreamConnection>);

#[cfg(feature = "mcp")]
impl bitrouter_core::tools::provider::ToolProvider for McpToolProviderAdapter {
    fn provider_name(&self) -> &str {
        bitrouter_core::tools::provider::ToolProvider::provider_name(self.0.as_ref())
    }

    async fn call_tool(
        &self,
        tool_id: &str,
        arguments: serde_json::Value,
    ) -> bitrouter_core::errors::Result<bitrouter_core::tools::result::ToolCallResult> {
        bitrouter_core::tools::provider::ToolProvider::call_tool(
            self.0.as_ref(),
            tool_id,
            arguments,
        )
        .await
    }
}

// ── REST auth helpers ────────────────────────────────────────────

/// Resolve the auth header from a provider config.
///
/// Priority: `auth` config (with resolved api_key) > `api_key` field (default
/// to Bearer). When the `auth.api_key` is an unsubstituted env var placeholder,
/// falls back to the provider-level `api_key` (which is resolved by
/// `env_prefix` during config loading).
#[cfg(feature = "rest")]
pub(crate) fn resolve_auth_header(config: &ProviderConfig) -> Option<(String, String)> {
    use bitrouter_config::AuthConfig;
    match config.auth.as_ref() {
        Some(AuthConfig::Bearer { api_key }) => {
            let key = resolve_key(api_key, config);
            Some(("Authorization".to_owned(), format!("Bearer {key}")))
        }
        Some(AuthConfig::Header {
            header_name,
            api_key,
        }) => {
            let key = resolve_key(api_key, config);
            Some((header_name.clone(), key))
        }
        Some(
            AuthConfig::X402
            | AuthConfig::Mpp
            | AuthConfig::Wallet
            | AuthConfig::Custom { .. }
            | AuthConfig::OAuth { .. },
        ) => None,
        None => {
            // Fall back to api_key field as Bearer token.
            config
                .api_key
                .as_ref()
                .map(|key| ("Authorization".to_owned(), format!("Bearer {key}")))
        }
    }
}

/// If the key is an unsubstituted env var placeholder (e.g. `"${EXA_API_KEY}"`),
/// fall back to the provider-level resolved `api_key`.
#[cfg(feature = "rest")]
fn resolve_key(auth_key: &str, config: &ProviderConfig) -> String {
    if auth_key.starts_with("${") && auth_key.ends_with('}') {
        config
            .api_key
            .clone()
            .unwrap_or_else(|| auth_key.to_owned())
    } else {
        auth_key.to_owned()
    }
}

// ── Tool call handler ────────────────────────────────────────────

/// [`ToolCallHandler`] implementation that dispatches `tools/call` through
/// a [`ToolRouter`] dispatch chain.
///
/// Routes tool calls using the config-authoritative routing table for
/// name resolution, then dispatches through the tool router (which may
/// wrap providers with guardrail enforcement via [`GuardedToolRouter`]).
#[cfg(feature = "mcp")]
pub struct RouterToolCallHandler<R, T> {
    tool_router: Arc<R>,
    tool_table: Arc<T>,
}

#[cfg(feature = "mcp")]
impl<R, T> RouterToolCallHandler<R, T> {
    pub fn new(tool_router: Arc<R>, tool_table: Arc<T>) -> Self {
        Self {
            tool_router,
            tool_table,
        }
    }
}

#[cfg(feature = "mcp")]
impl<R, T> bitrouter_core::api::mcp::gateway::ToolCallHandler for RouterToolCallHandler<R, T>
where
    R: bitrouter_core::routers::router::ToolRouter + Send + Sync + 'static,
    T: bitrouter_core::routers::routing_table::RoutingTable + Send + Sync + 'static,
{
    fn call_tool(
        &self,
        name: &str,
        arguments: Option<serde_json::Map<String, serde_json::Value>>,
    ) -> std::pin::Pin<
        Box<
            dyn std::future::Future<
                    Output = std::result::Result<
                        bitrouter_core::api::mcp::types::McpToolCallResult,
                        bitrouter_core::api::mcp::types::McpGatewayError,
                    >,
                > + Send
                + '_,
        >,
    > {
        let name = name.to_owned();
        Box::pin(async move {
            use bitrouter_core::api::mcp::types::McpGatewayError;
            use bitrouter_core::routers::content::RouteContext;
            use bitrouter_core::tools::provider::ToolProvider;

            let (_provider_name, tool_id) = name
                .split_once('/')
                .map(|(p, t)| (p.to_owned(), t.to_owned()))
                .unwrap_or_else(|| (name.clone(), name.clone()));

            // Route through config-authoritative table.
            let target = self
                .tool_table
                .route(&name, &RouteContext::default())
                .await
                .map_err(|e| McpGatewayError::ToolNotFound {
                    name: format!("{name}: {e}"),
                })?;

            // Dispatch through the tool router (GuardedToolRouter enforces
            // parameter restrictions on the returned provider).
            let provider_impl = self.tool_router.route_tool(target).await.map_err(|e| {
                McpGatewayError::UpstreamCall {
                    name: name.clone(),
                    reason: e.to_string(),
                }
            })?;

            let args_value = bitrouter_core::api::mcp::convert::args_to_value(arguments);
            let result = provider_impl
                .call_tool(&tool_id, args_value)
                .await
                .map_err(|e| McpGatewayError::UpstreamCall {
                    name,
                    reason: e.to_string(),
                })?;

            Ok(bitrouter_core::api::mcp::types::McpToolCallResult::from(
                result,
            ))
        })
    }
}

// ── Agent router ─────────────────────────────────────────────────

/// Config-driven agent router that creates and caches
/// [`AcpAgentProvider`] instances from the `agents:` config section.
///
/// Enabled agents get a shared provider on construction. Callers receive
/// an `Arc`-wrapped handle so that multiple API requests can share the
/// same session pool for a given agent.
#[cfg(feature = "tui")]
pub struct ConfigAgentRouter {
    agents: HashMap<String, Arc<bitrouter_providers::acp::provider::AcpAgentProvider>>,
}

#[cfg(feature = "tui")]
impl ConfigAgentRouter {
    /// Build a router from the `agents:` configuration map.
    ///
    /// Only enabled agents are instantiated. Disabled agents are silently
    /// skipped — they still appear in the discovery registry but cannot
    /// be routed to.
    pub fn new(agent_configs: HashMap<String, bitrouter_config::AgentConfig>) -> Self {
        let agents = agent_configs
            .into_iter()
            .filter(|(_, config)| config.enabled)
            .map(|(name, config)| {
                let provider = Arc::new(bitrouter_providers::acp::provider::AcpAgentProvider::new(
                    name.clone(),
                    config,
                ));
                (name, provider)
            })
            .collect();
        Self { agents }
    }

    /// Returns an iterator over the managed providers.
    ///
    /// Used by the runtime to drive idle-session cleanup across all
    /// agents.
    pub fn providers(
        &self,
    ) -> impl Iterator<Item = &Arc<bitrouter_providers::acp::provider::AcpAgentProvider>> {
        self.agents.values()
    }
}

#[cfg(feature = "tui")]
impl bitrouter_core::routers::router::AgentRouter for ConfigAgentRouter {
    async fn route_agent(
        &self,
        agent_name: &str,
    ) -> Result<Box<bitrouter_core::agents::provider::DynAgentProvider<'static>>> {
        let provider = self.agents.get(agent_name).ok_or_else(|| {
            BitrouterError::invalid_request(None, format!("unknown agent: {agent_name}"), None)
        })?;
        Ok(bitrouter_core::agents::provider::DynAgentProvider::new_box(
            Arc::clone(provider),
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bitrouter_core::routers::router::ToolRouter;
    use bitrouter_core::tools::provider::ToolProvider;

    fn test_providers() -> HashMap<String, ProviderConfig> {
        let mut p = HashMap::new();
        p.insert(
            "test-rest".into(),
            ProviderConfig {
                api_protocol: Some(ApiProtocol::Rest),
                api_base: Some("https://example.com/api".into()),
                api_key: Some("test-key".into()),
                ..Default::default()
            },
        );
        p
    }

    #[tokio::test]
    async fn route_to_rest_provider() {
        let router = LazyToolRouter::new(
            test_providers(),
            #[cfg(feature = "mcp")]
            Arc::new(HashMap::new()),
            Arc::new(reqwest::Client::new()),
        );

        let target = RoutingTarget {
            provider_name: "test-rest".into(),
            service_id: "search".into(),
            api_protocol: ApiProtocol::Rest,
        };

        let provider = router.route_tool(target).await;
        assert!(provider.is_ok());
        let p = provider.as_ref().ok();
        assert_eq!(p.map(|p| p.provider_name()), Some("test-rest"));
    }

    #[tokio::test]
    async fn route_to_unknown_provider_errors() {
        let router = LazyToolRouter::new(
            HashMap::new(),
            #[cfg(feature = "mcp")]
            Arc::new(HashMap::new()),
            Arc::new(reqwest::Client::new()),
        );

        let target = RoutingTarget {
            provider_name: "missing".into(),
            service_id: "foo".into(),
            api_protocol: ApiProtocol::Rest,
        };
        assert!(router.route_tool(target).await.is_err());
    }

    #[tokio::test]
    async fn has_providers_rest() {
        let router = LazyToolRouter::new(
            test_providers(),
            #[cfg(feature = "mcp")]
            Arc::new(HashMap::new()),
            Arc::new(reqwest::Client::new()),
        );
        assert!(router.has_providers());
    }

    #[tokio::test]
    async fn has_providers_empty() {
        let router = LazyToolRouter::new(
            HashMap::new(),
            #[cfg(feature = "mcp")]
            Arc::new(HashMap::new()),
            Arc::new(reqwest::Client::new()),
        );
        assert!(!router.has_providers());
    }

    // ── ConfigAgentRouter ────────────────────────────────────────────

    #[cfg(feature = "tui")]
    fn test_agent_configs() -> HashMap<String, bitrouter_config::AgentConfig> {
        let mut m = HashMap::new();
        m.insert(
            "claude-code".to_owned(),
            bitrouter_config::AgentConfig {
                protocol: bitrouter_config::AgentProtocol::Acp,
                binary: "claude".to_owned(),
                args: vec!["--agent".to_owned()],
                enabled: true,
                distribution: Vec::new(),
                session: Some(bitrouter_config::AgentSessionConfig {
                    idle_timeout_secs: 300,
                    max_concurrent: 5,
                }),
                a2a: None,
            },
        );
        m.insert(
            "disabled-agent".to_owned(),
            bitrouter_config::AgentConfig {
                protocol: bitrouter_config::AgentProtocol::Acp,
                binary: "disabled".to_owned(),
                args: Vec::new(),
                enabled: false,
                distribution: Vec::new(),
                session: None,
                a2a: None,
            },
        );
        m
    }

    #[cfg(feature = "tui")]
    #[tokio::test]
    async fn agent_router_routes_enabled_agent() {
        use bitrouter_core::routers::router::AgentRouter;
        let router = super::ConfigAgentRouter::new(test_agent_configs());
        let provider = router.route_agent("claude-code").await;
        assert!(provider.is_ok());
    }

    #[cfg(feature = "tui")]
    #[tokio::test]
    async fn agent_router_rejects_unknown_agent() {
        use bitrouter_core::routers::router::AgentRouter;
        let router = super::ConfigAgentRouter::new(test_agent_configs());
        assert!(router.route_agent("nonexistent").await.is_err());
    }

    #[cfg(feature = "tui")]
    #[tokio::test]
    async fn agent_router_skips_disabled_agent() {
        use bitrouter_core::routers::router::AgentRouter;
        let router = super::ConfigAgentRouter::new(test_agent_configs());
        // Disabled agents should not be routable.
        assert!(router.route_agent("disabled-agent").await.is_err());
    }

    #[cfg(feature = "tui")]
    #[tokio::test]
    async fn agent_router_providers_iterator() {
        let router = super::ConfigAgentRouter::new(test_agent_configs());
        // Only enabled agents are in the provider pool.
        let count = router.providers().count();
        assert_eq!(count, 1);
    }
}