use axum::extract::{
    State,
    ws::{Message, WebSocket, WebSocketUpgrade},
};
use dashmap::DashMap;
use r2d2::Pool;
use r2d2_sqlite::SqliteConnectionManager;
use std::sync::Arc;
use tokio::sync::mpsc;

use crate::utils::websocket::WsMap;

#[derive(Clone)]
pub struct AppState {
    pub hotel_pools: HotelPool,
    pub logs_pool: Pool<SqliteConnectionManager>,
    pub ws_map: WsMap,
}

type HotelId = i32;

#[derive(Clone)]
pub struct HotelPool {
    hotel_pools: Arc<DashMap<HotelId, Pool<SqliteConnectionManager>>>,
}

impl HotelPool {
    pub fn new() -> Self {
        Self {
            hotel_pools: Arc::new(DashMap::new()),
        }
    }

    pub fn get_pool(&self, hotel_id: i32) -> Pool<SqliteConnectionManager> {
        if let Some(pool) = self.hotel_pools.get(&hotel_id) {
            return pool.clone();
        }

        let db_path = format!("db/{}.sqlite", hotel_id);
        let hotels_manager = SqliteConnectionManager::file(db_path);
        let db_pool = Pool::builder()
            .max_size(5) // adjust based on load
            .build(hotels_manager)
            .expect("Failed to build pool");

        self.hotel_pools.insert(hotel_id, db_pool.clone());
        db_pool
    }
}
