pub mod catalog;
pub mod error;
pub mod settings;
pub mod space_token;

mod auth;
mod net;
mod server;
mod services;
mod tools;

pub use server::McpServer;
pub use services::{McpServices, SharedServices};

use std::net::IpAddr;
use std::sync::Arc;

use axum::Router;
use axum::middleware;
use export::AppState;
use ipnet::IpNet;
use rmcp::transport::StreamableHttpServerConfig;
use rmcp::transport::StreamableHttpService;
use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
use sha2::{Digest, Sha256};
use uuid::Uuid;

struct InternalKey {
    name: String,
    hash: String,
}

pub(crate) struct RuntimeInner {
    pub(crate) services: SharedServices,
    internal_keys: Vec<InternalKey>,
    internal_cidrs: Vec<IpNet>,
    trusted_proxies: Vec<IpNet>,
    system_uid: Option<i64>,
}

#[derive(Clone)]
pub(crate) struct McpRuntime(Arc<RuntimeInner>);

impl McpRuntime {
    pub(crate) fn services(&self) -> &SharedServices {
        &self.0.services
    }

    pub(crate) fn system_uid(&self) -> Option<i64> {
        self.0.system_uid
    }

    pub(crate) fn trusted_proxies(&self) -> &[IpNet] {
        &self.0.trusted_proxies
    }

    pub(crate) fn match_internal_key(&self, token: &str) -> Option<&str> {
        let hash = sha256_hex(token.as_bytes());
        self.0
            .internal_keys
            .iter()
            .find(|key| constant_time_eq(hash.as_bytes(), key.hash.as_bytes()))
            .map(|key| key.name.as_str())
    }

    pub(crate) fn is_internal_ip(&self, ip: Option<IpAddr>) -> bool {
        ip.is_some_and(|ip| self.0.internal_cidrs.iter().any(|net| net.contains(&ip)))
    }
}

pub async fn router(state: AppState) -> Router {
    let config = state.config.mcp.clone();
    if !config.enabled() {
        return Router::new();
    }

    let system_uid = match agent::bot::resolve_bot_user(&state, &state.config.ai)
        .await
    {
        Ok(uid) => {
            export::authz::set_system_uid(uid);
            Some(uid)
        }
        Err(error) => {
            tracing::warn!(%error, "failed to resolve the platform bot user for mcp internal keys");
            None
        }
    };

    let internal_keys = config
        .internal_keys
        .iter()
        .filter(|key| key.enabled())
        .map(|key| InternalKey {
            name: key.name.clone(),
            hash: key
                .token_hash
                .as_ref()
                .map(|hash| hash.trim().to_lowercase())
                .unwrap_or_else(|| sha256_hex(key.token.as_deref().unwrap_or_default().as_bytes())),
        })
        .collect();

    let runtime = McpRuntime(Arc::new(RuntimeInner {
        services: Arc::new(McpServices::new(state)),
        internal_keys,
        internal_cidrs: config.parsed_internal_cidrs(),
        trusted_proxies: config.parsed_trusted_proxies(),
        system_uid,
    }));

    let mut http_config = StreamableHttpServerConfig::default()
        .with_legacy_session_mode(false)
        .with_json_response(true);
    if !config.allowed_hosts.is_empty() {
        http_config = http_config.with_allowed_hosts(config.allowed_hosts.clone());
    }

    let services = runtime.services().clone();
    let mcp_service = StreamableHttpService::new(
        move || Ok(McpServer::new(services.clone())),
        LocalSessionManager::default().into(),
        http_config,
    );

    tracing::info!(
        internal_keys = runtime.0.internal_keys.len(),
        internal_cidrs = runtime.0.internal_cidrs.len(),
        "mcp endpoint ready"
    );

    Router::new()
        .fallback_service(mcp_service)
        .layer(middleware::from_fn_with_state(
            runtime,
            auth::auth_middleware,
        ))
}

pub(crate) fn sha256_hex(data: &[u8]) -> String {
    let digest = Sha256::digest(data);
    digest.iter().map(|byte| format!("{byte:02x}")).collect()
}

fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut diff: u8 = 0;
    for (x, y) in a.iter().zip(b.iter()) {
        diff |= x ^ y;
    }
    diff == 0
}