use std::collections::{HashMap, HashSet};
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use anyhow::Context;
use axum::extract::ws::Message;
use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
use tokio::sync::{mpsc, oneshot, Mutex, RwLock};
use tokio_postgres::NoTls;

use crate::logger::Logger;
use crate::types::{
    JsonValue, MicroserviceProxyCacheEntry, OperationPerformanceSample, RequestPerformanceSample,
    RuntimeConfig,
};

pub struct ProviderConnection {
    pub provider_id: Mutex<Option<String>>,
    pub sender: mpsc::UnboundedSender<Message>,
}

pub struct SshClientConnection {
    pub provider_id: String,
    pub sender: mpsc::UnboundedSender<Message>,
    pub data_channel_id: String,
}

pub struct SshDataFrame {
    pub header: JsonValue,
    pub payload: Vec<u8>,
}

pub struct SshDataChannel {
    pub provider_id: String,
    pub channel_id: String,
    pub writer: mpsc::UnboundedSender<SshDataFrame>,
    pub active_session_id: Option<String>,
    pub connected_at_millis: i64,
}

pub struct HttpTunnelResponse {
    pub provider_id: String,
    pub request_id: String,
    pub ok: bool,
    pub result: JsonValue,
}

pub type HttpTunnelWaiter = oneshot::Sender<Result<HttpTunnelResponse, String>>;
pub type TaskTerminalWaiter = oneshot::Sender<Option<crate::types::RawTaskRow>>;

pub struct AppState {
    pub config: RuntimeConfig,
    pub pool: Pool,
    pub logger: Logger,
    db_ready: AtomicBool,
    pub service_started_at: chrono::DateTime<chrono::Utc>,
    pub active_providers: RwLock<HashMap<String, Arc<ProviderConnection>>>,
    pub active_ssh_clients: RwLock<HashMap<String, SshClientConnection>>,
    pub active_ssh_data_channels: Mutex<HashMap<String, SshDataChannel>>,
    pub http_tunnel_waiters: Mutex<HashMap<String, HttpTunnelWaiter>>,
    pub task_terminal_waiters: Mutex<HashMap<String, Vec<TaskTerminalWaiter>>>,
    pub active_egress_tcp: Mutex<HashMap<String, crate::egress_tcp::EgressTcpConnection>>,
    pub microservice_proxy_cache: Mutex<HashMap<String, MicroserviceProxyCacheEntry>>,
    pub microservice_availability_cache: Mutex<HashMap<String, (i64, JsonValue)>>,
    pub request_performance_samples: Mutex<Vec<RequestPerformanceSample>>,
    pub operation_performance_samples: Mutex<Vec<OperationPerformanceSample>>,
    pub active_scheduled_runs: Mutex<std::collections::HashSet<String>>,
    pub active_provider_upgrade_watchers: Mutex<HashSet<String>>,
}

impl AppState {
    pub async fn new(config: RuntimeConfig, logger: Logger) -> anyhow::Result<Self> {
        let pg_config =
            tokio_postgres::Config::from_str(&config.database_url).context("parse DATABASE_URL")?;
        let manager = Manager::from_config(
            pg_config,
            NoTls,
            ManagerConfig {
                recycling_method: RecyclingMethod::Fast,
            },
        );
        let pool = Pool::builder(manager)
            .max_size(config.database_pool_max)
            .build()
            .context("build PostgreSQL pool")?;
        Ok(Self {
            config,
            pool,
            logger,
            db_ready: AtomicBool::new(false),
            service_started_at: chrono::Utc::now(),
            active_providers: RwLock::new(HashMap::new()),
            active_ssh_clients: RwLock::new(HashMap::new()),
            active_ssh_data_channels: Mutex::new(HashMap::new()),
            http_tunnel_waiters: Mutex::new(HashMap::new()),
            task_terminal_waiters: Mutex::new(HashMap::new()),
            active_egress_tcp: Mutex::new(HashMap::new()),
            microservice_proxy_cache: Mutex::new(HashMap::new()),
            microservice_availability_cache: Mutex::new(HashMap::new()),
            request_performance_samples: Mutex::new(Vec::new()),
            operation_performance_samples: Mutex::new(Vec::new()),
            active_scheduled_runs: Mutex::new(std::collections::HashSet::new()),
            active_provider_upgrade_watchers: Mutex::new(HashSet::new()),
        })
    }

    pub fn set_db_ready(&self, ready: bool) {
        self.db_ready.store(ready, Ordering::SeqCst);
    }

    pub fn db_ready(&self) -> bool {
        self.db_ready.load(Ordering::SeqCst)
    }

    pub fn log(&self, level: &str, message: &str, data: Option<JsonValue>) {
        self.logger.log(level, message, data);
    }

    pub async fn active_provider_count(&self) -> usize {
        self.active_providers.read().await.len()
    }

    pub async fn send_provider(&self, provider_id: &str, value: JsonValue) -> bool {
        let providers = self.active_providers.read().await;
        providers.get(provider_id).is_some_and(|provider| {
            provider
                .sender
                .send(Message::Text(value.to_string()))
                .is_ok()
        })
    }
}
