pub use crate::listener::{RequestHandler, RequestReceiver};

use crate::pool::Pool;

pub use crate::pool::RequestSender;

use lyquor_primitives::NodeID;

use lyquor_tls::TlsConfig;

use std::collections::HashMap;

use std::sync::Arc;

use thiserror::Error;

#[derive(Debug, Error)]

pub enum HubError {

    #[error("invalid address")]

    InvalidAddr,

    #[error("peer not exist")]

    PeerNotExist,

    #[error("TLS error: {0}")]

    TlsError(#[from] lyquor_tls::TlsError),

    #[error("unknown error")]

    Unknown,

}

type Result<T> = std::result::Result<T, HubError>;

pub struct Hub {

    id: NodeID,

    peers: HashMap<NodeID, crate::pool::Pool>,

    listener: crate::listener::Listener,

    runtime: tokio::runtime::Handle,

    spawner: crate::event::AsyncSpawner,

    tls_config: TlsConfig,

}

impl std::fmt::Debug for Hub {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.debug_struct("Hub").field("id", &self.id).finish_non_exhaustive()

    }

}

impl Hub {

    #[tracing::instrument(level = "trace", skip(tls_config))]

    fn new(id: NodeID, listen_addr: String, tls_config: TlsConfig) -> Result<Self> {

        let runtime = tokio::runtime::Handle::current(); // TODO: use our runtime

        let server_config = tls_config.server_config().map_err(|e| HubError::TlsError(
e0
))
?0
;

        let listener = crate::listener::Listener::new(id, server_config, listen_addr);

        Ok(Hub {

            id,

            peers: HashMap::new(),

            listener: listener,

            runtime: runtime.clone(),

            spawner: crate::event::AsyncSpawner { handle: runtime },

            tls_config: tls_config,

        })

    }

    pub fn spawner(&self) -> crate::event::AsyncSpawner {

        self.spawner.clone()

    }

    pub fn get_id(&self) -> NodeID {

        self.id.clone()

    }

    pub fn signing_key(&self) -> Arc<lyquor_tls::NodeSigningKey> {

        self.tls_config.signing_key()

    }

    pub async fn start(&mut self) -> Result<()> {

        match self.listener.start().await {

            // listener.start() only returns bind error

            Err(_) => Err(HubError::InvalidAddr),

            Ok(_) => Ok(()),

        }

    }

    #[tracing::instrument(level = "trace", ret, err)]

    pub async fn add_peer(&mut self, id: NodeID, addr: String) -> Result<()> {

        tracing::info!("add peer for {:?}: {:?} -> {:?}", self.id, id, addr);

        self.listener.add_peer(id.clone()).await;

        let config = crate::pool::PoolConfig::builder()

            .tls_config(self.tls_config.client_config().unwrap())

            .addr(addr)

            .build();

        let pool = Pool::new(config);

        let _ = pool.start(self.runtime.clone()).await;

        // TODO: make sure the pool connected to the node matches the id.

        self.peers.insert(id, pool);

        Ok(())

    }

    #[tracing::instrument(level = "trace", ret, err(level = "debug"))]

    pub async fn remove_peer(&mut self, id: NodeID) -> Result<()> {

        self.peers.remove(&id);

        Ok(())

    }

    #[tracing::instrument(level = "trace", err(level = "debug"))]

    pub fn outbound(&self, id: NodeID) -> Result<RequestSender> {

        let pool = self.peers.get(&id).ok_or_else(|| HubError::PeerNotExist)?;

        Ok(RequestSender::new(self.id.clone(), pool))

    }

    #[tracing::instrument(level = "trace", err(level = "debug"))]

    pub fn inbound(&self, id: NodeID) -> Result<RequestReceiver> {

        self.listener.inbound(id).map_err(|_e| HubError::Unknown)

    }

}

impl Drop for Hub {

    fn drop(&mut self) {

        self.listener.stop();

    }

}

pub struct HubBuilder {

    tls_config: TlsConfig,

    listen_addr: String,

}

impl HubBuilder {

    pub fn new(tls_config: TlsConfig) -> Self {

        HubBuilder {

            tls_config: tls_config,

            listen_addr: "127.0.0.1:0".to_string(),

        }

    }

    pub fn listen_addr(mut self, addr: String) -> Self {

        self.listen_addr = addr;

        self

    }

    pub fn build(self) -> Result<Hub> {

        Hub::new(self.tls_config.node_id(), self.listen_addr, self.tls_config)

    }

}

#[cfg(test)]

pub(crate) mod tests {

    use super::*;

    use crate::listener::ResponseBuilder;

    use bytes::Bytes;

    use http_body_util::BodyExt;

    use lyquor_test::test;

    use std::sync::Arc;

    use tokio::sync::Mutex;

    fn generate_unused_socket() -> String {

        std::net::TcpListener::bind("127.0.0.1:0")

            .unwrap()

            .local_addr()

            .unwrap()

            .to_string()

    }

    #[test(tokio::test)]

    async fn test_hub() {

        let mut nodes = Vec::new();

        for i in 0u8..2 {

            let (node_id, tls_config) = lyquor_tls::generator::test_config(i);

            let addr = generate_unused_socket();

            let mut hub = HubBuilder::new(tls_config).listen_addr(addr.clone()).build().unwrap();

            let _ = hub.start().await;

            nodes.push((node_id, addr, Arc::new(Mutex::new(hub))));

        }

        for (id, _, hub) in &nodes {

            for (peer_id, peer_addr, _) in &nodes {

                if id == peer_id {

                    continue;

                }

                let mut hub = hub.lock().await;

                hub.add_peer(peer_id.clone(), peer_addr.clone()).await.unwrap();

            }

        }

        let first = nodes[0].2.lock().await;

        let second = nodes[1].2.lock().await;

        let requester = first.outbound(nodes[1].0.clone()).unwrap();

        let responder = second.inbound(nodes[0].0.clone()).unwrap();

        tokio::spawn(async move {

            loop {

                let 
req1
 = responder.recv().await.
unwrap1
();

                let req_body = req.request.collect().await.unwrap().to_bytes();

                assert_eq!(req_body, Bytes::from("Hello, World!"));

                let res = ResponseBuilder::response("Hello, World!");

                req.response.send(res).unwrap();

            }

        });

        let response = requester

            .send_request(hyper::Request::new(http_body_util::Full::new(Bytes::from(

                "Hello, World!",

            ))))

            .await;

        dbg!(&response);

        assert!(response.is_ok());

        let response = response.unwrap();

        assert_eq!(response.status(), hyper::StatusCode::OK);

        let resp_body = response.collect().await.unwrap().to_bytes();

        assert_eq!(resp_body, Bytes::from("Hello, World!"));

    }

    #[test(tokio::test)]

    async fn test_cc() {

        use async_trait::async_trait;

        use hello::hello_client::HelloClient;

        use hello::hello_server::{Hello, HelloServer};

        use hello::{Command, Empty, HelloMsg, RespMsg};

        use rand::Rng;

        use rand::RngCore;

        use lyquor_net_rpc::TowerService;

        pub mod hello {

            lyquor_net_rpc::include_proto!("lyquor.hello");

        }

        fn gen_hellomsg(pld: Bytes) -> HelloMsg {

            HelloMsg {

                v_int32: 1,

                v_int64: 2,

                v_uint32: 3,

                v_uint64: 4,

                v_sint32: 5,

                v_sint64: 6,

                v_fixed32: 7,

                v_fixed64: 8,

                v_sfixed32: 9,

                v_sfixed64: 10,

                v_float: 11.0,

                v_double: 12.0,

                v_bool: true,

                v_string: "HelloMsg".to_string(),

                v_bytes: [1, 2, 3].to_vec(),

                cmd: Command::Pong.into(),

                len: pld.len() as u32,

                payload: pld.to_vec(),

            }

        }

        #[derive(Debug, Default)]

        pub struct HelloImpl {}

        #[async_trait]

        impl Hello for HelloImpl {

            async fn hello(

                &self, request: hyper::Request<HelloMsg>,

            ) -> std::result::Result<hyper::Response<HelloMsg>, String> {

                let (_, msg) = request.into_parts();

                println!("hello got a request: {:?}", msg);

                let mut rng = rand::rng();

                let mut pld = vec![0u8; rand::rng().random_range(100..=1000)];

                rng.fill_bytes(&mut pld);

                Ok(hyper::Response::new(gen_hellomsg(pld.into())))

            }

            async fn hello_empty(

                &self, request: hyper::Request<HelloMsg>,

            ) -> std::result::Result<hyper::Response<Empty>, String> {

                let (_, msg) = request.into_parts();

                println!("hello_empty got a request: {:?}", msg);

                let reply = Empty {};

                Ok(hyper::Response::new(reply))

            }

            async fn hello_resp(

                &self, request: hyper::Request<HelloMsg>,

            ) -> std::result::Result<hyper::Response<RespMsg>, String> {

                let (_, msg) = request.into_parts();

                println!("hello_resp got a request: {:?}", msg);

                let mut rng = rand::rng();

                let mut pld = vec![0u8; rand::rng().random_range(10..=20)];

                rng.fill_bytes(&mut pld);

                let reply = RespMsg {

                    cmd: Command::Pong.into(),

                    len: pld.len() as u32,

                    payload: pld,

                };

                Ok(hyper::Response::new(reply))

            }

            async fn hello_long(

                &self, request: hyper::Request<HelloMsg>,

            ) -> std::result::Result<hyper::Response<RespMsg>, String> {

                let (_, msg) = request.into_parts();

                println!("hello_long got a request: {:?}", msg);

                let mut rng = rand::rng();

                // 10M - 100M payload

                let mut pld = vec![0u8; rand::rng().random_range(10 * 1024 * 1024..=100 * 1024 * 1024)];

                rng.fill_bytes(&mut pld);

                let reply = RespMsg {

                    cmd: Command::Pong.into(),

                    len: pld.len() as u32,

                    payload: pld,

                };

                Ok(hyper::Response::new(reply))

            }

            async fn hello_error(

                &self, _: hyper::Request<HelloMsg>,

            ) -> std::result::Result<hyper::Response<RespMsg>, String> {

                Err("Error!".to_string())

            }

        }

        let mut nodes = Vec::new();

        for i in 0u8..2 {

            let (node_id, tls_config) = lyquor_tls::generator::test_config(i);

            let addr = generate_unused_socket();

            let mut hub = HubBuilder::new(tls_config).listen_addr(addr.clone()).build().unwrap();

            let _ = hub.start().await;

            nodes.push((node_id, addr, Arc::new(Mutex::new(hub))));

        }

        for (id, _, hub) in &nodes {

            for (peer_id, peer_addr, _) in &nodes {

                if id == peer_id {

                    continue;

                }

                let mut hub = hub.lock().await;

                hub.add_peer(peer_id.clone(), peer_addr.clone()).await.unwrap();

            }

        }

        let first = nodes[0].2.lock().await;

        let second = nodes[1].2.lock().await;

        let requester = first.outbound(nodes[1].0.clone()).unwrap();

        let responder = second.inbound(nodes[0].0.clone()).unwrap();

        let mut hello_srv = HelloServer::new(HelloImpl::default());

        tokio::spawn(async move {

            loop {

                let 
req5
 = responder.recv().await.
unwrap5
();

                println!("{:?}", req.request);

                assert_eq!(req.request.headers().get("lyquor-pkt-type").unwrap(), "RPC");

                let res = hello_srv.call(req.request).await.unwrap();

                req.response.send(res).unwrap();

            }

        });

        let msg = gen_hellomsg([0].to_vec().into());

        let mut client = HelloClient::new(requester);

        let resp = client.hello(msg.clone()).await;

        let response = resp.unwrap();

        assert_eq!(response.len, response.payload.len() as u32);

        println!("hello RESP: {:?}", response);

        let resp = client.hello_empty(msg.clone()).await;

        let response = resp.unwrap();

        println!("hello_empty RESP: {:?}", response);

        let resp = client.hello_resp(msg.clone()).await;

        let response = resp.unwrap();

        assert_eq!(response.len, response.payload.len() as u32);

        println!("hello_resp RESP: {:?}", response);

        let resp = client.hello_long(msg.clone()).await;

        let response = resp.unwrap();

        assert_eq!(response.len, response.payload.len() as u32);

        println!("hello_resp RESP length {:?}KB", response.len / 1023);

        let resp = client.hello_error(msg.clone()).await;

        match resp {

            Ok(_) => {

                // Shouldn't see this

                assert_eq!(0, 1);

            }

            Err(e) => {

                println!("{:?}", e);

            }

        };

    }

}