use std::{path::Path, sync::Arc};

use compio::net::{TcpListener, ToSocketAddrsAsync, UnixListener};
use compio::runtime::RuntimeBuilder;

use super::{
    handler::{KsHandler, RpcHandler, TxKsHandler},
    server::*,
};
use crate::logdb::db::Db;
use crate::transactional::db::TxDb;

impl Db {
    fn build_tree_map(&self) -> TreeMap {
        let db_info = self.db_info.cloned();
        let keyspaces = self.get_all_keyspaces();
        Arc::new(
            keyspaces
                .into_iter()
                .map(|(name, tree)| {
                    let hashname = xxhash_rust::xxh3::xxh3_64(name.as_bytes());
                    let info = db_info.collections.get(&name).copied().unwrap_or_default();
                    let handler = KsHandler {
                        name,
                        info,
                        tree,
                        seq_tree: self.seq_tree.clone(),
                    };
                    (hashname, Arc::new(handler) as Arc<dyn RpcHandler>)
                })
                .collect(),
        )
    }

    pub fn listen_tcp(&self, addr: impl ToSocketAddrsAsync + Send + 'static) {
        let trees: TreeMap = self.build_tree_map();
        std::thread::spawn(move || {
            RuntimeBuilder::new()
                .build()
                .expect("Failed to build compio runtime")
                .block_on(async {
                    let listener = TcpListener::bind(addr)
                        .await
                        .expect("Failed to bind TCP listener");
                    accept_tcp(listener, trees).await;
                });
        });
    }

    pub fn listen_uds(&self, path: impl AsRef<Path> + Send + 'static) {
        let trees: TreeMap = self.build_tree_map();
        std::thread::spawn(move || {
            RuntimeBuilder::new()
                .build()
                .expect("Failed to build compio runtime")
                .block_on(async {
                    let _ = std::fs::remove_file(path.as_ref());
                    let listener = UnixListener::bind(path)
                        .await
                        .expect("Failed to bind UDS listener");
                    accept_uds(listener, trees).await;
                });
        });
    }
}

impl TxDb {
    fn build_tree_map(&self) -> TreeMap {
        let db_info = self.db_info.cloned();
        let keyspaces = self.get_all_keyspaces();
        Arc::new(
            keyspaces
                .into_iter()
                .map(|(name, tree)| {
                    let hashname = xxhash_rust::xxh3::xxh3_64(name.as_bytes());
                    let info = db_info.collections.get(&name).copied().unwrap_or_default();
                    let handler = TxKsHandler {
                        name,
                        info,
                        tree,
                        seq_tree: self.seq_tree.clone(),
                        db: self.db.clone(),
                    };
                    (hashname, Arc::new(handler) as Arc<dyn RpcHandler>)
                })
                .collect(),
        )
    }

    pub fn listen_tcp(&self, addr: impl ToSocketAddrsAsync + Send + 'static) {
        let trees: TreeMap = self.build_tree_map();
        std::thread::spawn(move || {
            RuntimeBuilder::new()
                .build()
                .expect("Failed to build compio runtime")
                .block_on(async {
                    let listener = TcpListener::bind(addr)
                        .await
                        .expect("Failed to bind TCP listener");
                    accept_tcp(listener, trees).await;
                });
        });
    }

    pub fn listen_uds(&self, path: impl AsRef<Path> + Send + 'static) {
        let trees: TreeMap = self.build_tree_map();
        std::thread::spawn(move || {
            RuntimeBuilder::new()
                .build()
                .expect("Failed to build compio runtime")
                .block_on(async {
                    let _ = std::fs::remove_file(path.as_ref());
                    let listener = UnixListener::bind(path)
                        .await
                        .expect("Failed to bind UDS listener");
                    accept_uds(listener, trees).await;
                });
        });
    }
}
