use net_pool::{debug, info, instrument_debug_span, tokio_spawn, warn2};
use net_relay::{Builder, Error};
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;

/// tcp relay
pub struct Relay<F, S, P = tcp_pool::Pool>
where
    F: Fn(Arc<P>, tokio::net::TcpStream, SocketAddr) -> S,
    S: Future<Output = ()>,
    P: tcp_pool::TcpPool + net_pool::Pool,
{
    parts: net_relay::builder::Parts<P, F>,
    pending: Option<Pin<Box<dyn Future<Output = Result<(), net_relay::Error>> + Send + 'static>>>,
}

impl<F, S, P> Relay<F, S, P>
where
    F: Fn(Arc<P>, tokio::net::TcpStream, SocketAddr) -> S,
    S: Future<Output = ()>,
    P: tcp_pool::TcpPool + net_pool::Pool,
{
    pub fn build<B: FnOnce(Builder<P, F>) -> Builder<P, F>>(b: B) -> Result<Self, Error> {
        let builder = Builder::new();
        let parts = b(builder).build()?;
        Ok(Relay {
            parts,
            pending: None,
        })
    }

    pub fn bind_addrs(&self) -> &Vec<SocketAddr> {
        &self.parts.bind_addrs
    }

    pub fn relay_fn(&self) -> Arc<F> {
        self.parts.relay_fn.as_ref().unwrap().clone()
    }

    pub fn pool(&self) -> Arc<P> {
        self.parts.pools[0].clone()
    }

    /// 设置最大连接数
    pub fn set_max_conn(&self, max: Option<usize>) {
        self.pool().set_max_conn(max)
    }

    /// 设置空闲连接保留时长
    pub fn set_keepalive(&self, duration: Option<Duration>) {
        self.pool().set_keepalive(duration)
    }

    /// 添加一个后端地址
    pub fn add_backend(&self, addr: net_pool::backend::Address) {
        self.pool().add_backend(addr)
    }

    /// 移除一个后端地址
    pub fn remove_backend(&self, addr: &net_pool::backend::Address) -> bool {
        self.pool().remove_backend(addr)
    }
}

impl<F, S, P> net_relay::Relay for Relay<F, S, P>
where
    F: Fn(Arc<P>, tokio::net::TcpStream, SocketAddr) -> S + Send + Sync + 'static,
    S: Future<Output = ()> + Send + 'static,
    P: tcp_pool::TcpPool + net_pool::Pool + Send + 'static,
{
    fn poll_run(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        if self.pending.is_none() {
            let tuple = (self.bind_addrs().clone(), self.pool(), self.relay_fn());

            self.pending = Some(Box::pin(async move {
                let listener = tokio::net::TcpListener::bind(tuple.0.as_slice()).await?;

                info!(
                    "[tcp service] listen on: {:?}",
                    listener.local_addr().unwrap()
                );

                loop {
                    match listener.accept().await {
                        Ok((client, addr)) => {
                            let tuple = (tuple.1.clone(), tuple.2.clone());
                            tokio_spawn! {
                                instrument_debug_span! {
                                    async move {
                                        debug!("[tcp service] connection accepted");
                                        let res = tuple.1(tuple.0, client, addr).await;
                                        debug!("[tcp service] connection closed");
                                        res
                                    },
                                    "new_tcp_stream",
                                    address=addr.to_string()
                                }
                            };
                        }
                        Err(_e) => {
                            warn2!("[Tcp Relay] accept from listen, error occurred: {:?}", _e);
                        }
                    }
                }
            }));
        }

        self.pending.as_mut().unwrap().as_mut().poll(cx)
    }
}

pub async fn default_relay_fn<P: net_pool::Pool + tcp_pool::TcpPool + Send>(
    pool: Arc<P>,
    mut client: tokio::net::TcpStream,
    addr: SocketAddr,
) {
    let id = net_pool::utils::socketaddr_to_hash_code(&addr);
    let mut proxy = match pool.clone().get(&id.to_string()).await {
        Err(_e) => {
            warn2!(
                "[Tcp Relay] get tcp stream from pool, error occurred: {:?}",
                _e
            );
            return;
        }
        Ok(t) => t,
    };

    let proxy_mut: &mut tokio::net::TcpStream = &mut proxy;
    let _cnt = tokio::io::copy_bidirectional(proxy_mut, &mut client).await;

    debug!("[Tcp Relay] data exchange byte count: {:?}", _cnt);
}
