use crate::proxy::TokenManager;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Clone)]
pub struct CoreServices {
    pub token_manager: Arc<TokenManager>,
    pub upstream: Arc<crate::proxy::upstream::client::UpstreamClient>,
    pub monitor: Arc<crate::proxy::monitor::ProxyMonitor>,
    pub integration: crate::modules::system::integration::SystemManager,
    pub account_service: Arc<crate::modules::auth::account_service::AccountService>,
}

#[derive(Clone)]
pub struct ConfigState {
    pub custom_mapping: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
    pub upstream_proxy: Arc<tokio::sync::RwLock<crate::proxy::config::UpstreamProxyConfig>>,
    pub zai: Arc<RwLock<crate::proxy::ZaiConfig>>,
    pub experimental: Arc<RwLock<crate::proxy::config::ExperimentalConfig>>,
    pub debug_logging: Arc<RwLock<crate::proxy::config::DebugLoggingConfig>>,
    pub security: Arc<RwLock<crate::proxy::ProxySecurityConfig>>,
    pub request_timeout: Arc<AtomicU64>,
    pub update_lock: Arc<tokio::sync::Mutex<()>>,
}

#[derive(Clone)]
pub struct RuntimeState {
    pub provider_rr: Arc<AtomicUsize>,
    pub switching: Arc<RwLock<bool>>,
    pub is_running: Arc<RwLock<bool>>,
    pub port: u16,
    pub proxy_pool_state: Arc<tokio::sync::RwLock<crate::proxy::config::ProxyPoolConfig>>,
    pub proxy_pool_manager: Arc<crate::proxy::proxy_pool::ProxyPoolManager>,
}
#[derive(Clone)]
pub struct AppState {
    pub core: Arc<CoreServices>,
    pub config: Arc<ConfigState>,
    pub runtime: Arc<RuntimeState>,
}

#[derive(Clone)]
pub struct AdminState {
    pub core: Arc<CoreServices>,
    pub config: Arc<ConfigState>,
    pub runtime: Arc<RuntimeState>,
}

#[derive(Clone)]
pub struct OpenAIHandlerState {
    pub token_manager: Arc<TokenManager>,
    pub custom_mapping: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
    pub upstream: Arc<crate::proxy::upstream::client::UpstreamClient>,
    pub debug_logging: Arc<RwLock<crate::proxy::config::DebugLoggingConfig>>,
}

#[derive(Clone)]
pub struct ModelCatalogState {
    pub custom_mapping: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
}

impl ConfigState {
    pub async fn apply_proxy_config(&self, proxy: &crate::proxy::config::ProxyConfig) {
        let _guard = self.update_lock.lock().await;

        {
            let mut mapping = self.custom_mapping.write().await;
            *mapping = proxy.custom_mapping.clone();
        }
        {
            let mut upstream_proxy = self.upstream_proxy.write().await;
            *upstream_proxy = proxy.upstream_proxy.clone();
        }
        {
            let mut security = self.security.write().await;
            *security = crate::proxy::ProxySecurityConfig::from_proxy_config(proxy);
        }
        crate::proxy::middleware::client_ip::set_trusted_proxies(proxy.trusted_proxies.clone());
        {
            let mut zai = self.zai.write().await;
            *zai = proxy.zai.clone();
        }
        {
            let mut experimental = self.experimental.write().await;
            *experimental = proxy.experimental.clone();
        }
        {
            let mut debug_logging = self.debug_logging.write().await;
            *debug_logging = proxy.debug_logging.clone();
        }
        self.request_timeout
            .store(proxy.request_timeout, Ordering::Relaxed);
    }

    pub fn request_timeout_secs(&self) -> u64 {
        self.request_timeout.load(Ordering::Relaxed)
    }
}

impl axum::extract::FromRef<AppState> for Arc<RwLock<crate::proxy::ProxySecurityConfig>> {
    fn from_ref(state: &AppState) -> Self {
        state.config.security.clone()
    }
}

impl axum::extract::FromRef<AppState> for CoreServices {
    fn from_ref(state: &AppState) -> Self {
        state.core.as_ref().clone()
    }
}

impl axum::extract::FromRef<AppState> for ConfigState {
    fn from_ref(state: &AppState) -> Self {
        state.config.as_ref().clone()
    }
}

impl axum::extract::FromRef<AppState> for RuntimeState {
    fn from_ref(state: &AppState) -> Self {
        state.runtime.as_ref().clone()
    }
}

impl axum::extract::FromRef<AppState> for AdminState {
    fn from_ref(state: &AppState) -> Self {
        Self {
            core: state.core.clone(),
            config: state.config.clone(),
            runtime: state.runtime.clone(),
        }
    }
}

impl axum::extract::FromRef<AppState> for Arc<CoreServices> {
    fn from_ref(state: &AppState) -> Self {
        state.core.clone()
    }
}

impl axum::extract::FromRef<AppState> for Arc<ConfigState> {
    fn from_ref(state: &AppState) -> Self {
        state.config.clone()
    }
}

impl axum::extract::FromRef<AppState> for Arc<RuntimeState> {
    fn from_ref(state: &AppState) -> Self {
        state.runtime.clone()
    }
}

impl axum::extract::FromRef<AppState> for OpenAIHandlerState {
    fn from_ref(state: &AppState) -> Self {
        Self {
            token_manager: state.core.token_manager.clone(),
            custom_mapping: state.config.custom_mapping.clone(),
            upstream: state.core.upstream.clone(),
            debug_logging: state.config.debug_logging.clone(),
        }
    }
}

impl axum::extract::FromRef<AppState> for ModelCatalogState {
    fn from_ref(state: &AppState) -> Self {
        Self {
            custom_mapping: state.config.custom_mapping.clone(),
        }
    }
}
