use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::Context;
use async_trait::async_trait;
use log::{debug, info};
use socks5_proto::Address;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
use tokio::net::TcpStream;
use tokio::time::timeout;

use crate::address_list::{DirectList, ProxyList};
use crate::proto::padding::Padding;
use crate::proto::trojan;
use crate::proto::trojan::Command;
use crate::tls::make_server_name;
use crate::tls::make_tls_connector;

pub trait AsyncStream: AsyncRead + AsyncWrite + Unpin + Send {}
impl<T: AsyncRead + AsyncWrite + Unpin + Send> AsyncStream for T {}

#[async_trait]
pub trait Dial: Send + Sync {
    async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>>;
}

pub struct DirectDial {
    connect_timeout: Duration,
}

impl DirectDial {
    pub fn new(connect_timeout: Duration) -> Self {
        Self { connect_timeout }
    }
}

#[async_trait]
impl Dial for DirectDial {
    async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
        let stream: TcpStream = match addr {
            Address::DomainAddress(domain, port) => {
                let domain = String::from_utf8_lossy(&domain);
                timeout(
                    self.connect_timeout,
                    TcpStream::connect((domain.as_ref(), port)),
                )
                .await
                .context(format!("connect {}:{} timeout", domain, port))?
                .context(format!("connect {}:{} failed", domain, port))
            }
            Address::SocketAddress(socket_addr) => {
                timeout(self.connect_timeout, TcpStream::connect(socket_addr))
                    .await
                    .context(format!("connect {} timeout", socket_addr))?
                    .context(format!("connect {} failed", socket_addr))
            }
        }?;
        Ok(Box::new(stream))
    }
}

pub struct TrojanDial {
    remote_addr: String,
    hash: String,
    insecure: bool,
    padding: bool,
    connect_timeout: Duration,
}

impl TrojanDial {
    pub fn new(
        remote_addr: String,
        hash: String,
        insecure: bool,
        padding: bool,
        connect_timeout: Duration,
    ) -> Self {
        Self {
            remote_addr,
            hash,
            insecure,
            padding,
            connect_timeout,
        }
    }
}

#[async_trait]
impl Dial for TrojanDial {
    async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
        let remote_ts = timeout(self.connect_timeout, TcpStream::connect(&self.remote_addr))
            .await
            .context(format!("connect {} timeout", self.remote_addr))?
            .context(format!("connect {} failed", self.remote_addr))?;

        let server_name = make_server_name(self.remote_addr.as_str())?;
        let mut remote_ts_ssl = make_tls_connector(self.insecure)
            .connect(server_name, remote_ts)
            .await
            .context("trojan can't connect tls")?;

        if self.padding {
            let req = trojan::Request::new(self.hash.clone(), Command::Padding, addr);
            req.write_to(&mut remote_ts_ssl).await?;
            Padding::read_from(&mut remote_ts_ssl).await?;
        } else {
            let req = trojan::Request::new(self.hash.clone(), Command::Connect, addr);
            req.write_to(&mut remote_ts_ssl).await?;
        }

        Ok(Box::new(remote_ts_ssl))
    }
}

#[cfg(feature = "websocket")]
pub struct WebSocketDial {
    remote_addr: String,
    hash: String,
    insecure: bool,
    padding: bool,
    connect_timeout: Duration,
}

#[cfg(feature = "websocket")]
impl WebSocketDial {
    pub fn new(
        remote_addr: String,
        hash: String,
        insecure: bool,
        padding: bool,
        connect_timeout: Duration,
    ) -> Self {
        Self {
            remote_addr,
            hash,
            insecure,
            padding,
            connect_timeout,
        }
    }
}

#[cfg(feature = "websocket")]
#[async_trait]
impl Dial for WebSocketDial {
    async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
        use crate::stream::websocket::WebSocketCopyStream;
        use crate::tls::make_tls_client_config;
        use bytes::BytesMut;
        use futures::SinkExt;
        use futures::StreamExt;
        use tokio_tungstenite::connect_async_tls_with_config;
        use tokio_tungstenite::tungstenite::Message;

        let (mut ws, _) = timeout(
            self.connect_timeout,
            connect_async_tls_with_config(
                &self.remote_addr,
                None,
                false,
                Some(tokio_tungstenite::Connector::Rustls(Arc::new(
                    make_tls_client_config(self.insecure),
                ))),
            ),
        )
        .await
        .context(format!("websocket connect {} timeout", self.remote_addr))?
        .context(format!("websocket connect {} failed", self.remote_addr))?;

        if self.padding {
            let mut buf = BytesMut::new();
            let req = trojan::Request::new(self.hash.clone(), Command::Padding, addr);
            req.write_to_buf(&mut buf);
            ws.send(Message::Binary(buf.freeze()))
                .await
                .context("websocket can't send")?;
            ws.flush().await?;
            let _ = ws.next().await;
        } else {
            let mut buf = BytesMut::new();
            let req = trojan::Request::new(self.hash.clone(), Command::Connect, addr);
            req.write_to_buf(&mut buf);
            ws.send(Message::Binary(buf.freeze()))
                .await
                .context("websocket can't send")?;
            ws.flush().await?;
        }
        Ok(Box::new(WebSocketCopyStream::new(ws)))
    }
}

const EXTRA_TIMEOUT_MS: u64 = 200;

pub struct SmartDial {
    direct: Box<dyn Dial>,
    proxy: Box<dyn Dial>,
    proxy_list: Arc<ProxyList>,
    direct_list: Arc<DirectList>,
    connect_timeout: Duration,
}

impl SmartDial {
    pub fn new(
        direct: Box<dyn Dial>,
        proxy: Box<dyn Dial>,
        proxy_list: Arc<ProxyList>,
        direct_list: Arc<DirectList>,
        connect_timeout: Duration,
    ) -> Self {
        Self {
            direct,
            proxy,
            proxy_list,
            direct_list,
            connect_timeout,
        }
    }

    async fn handle_proxy_result(
        &self,
        proxy_res: anyhow::Result<Box<dyn AsyncStream>>,
        addr: &Address,
        elapsed: Duration,
    ) -> anyhow::Result<Box<dyn AsyncStream>> {
        match proxy_res {
            Ok(mut stream) => {
                let remaining = self
                    .connect_timeout
                    .saturating_add(Duration::from_millis(EXTRA_TIMEOUT_MS))
                    .saturating_sub(elapsed);
                let check_timeout = remaining.min(Duration::from_millis(EXTRA_TIMEOUT_MS));

                if !is_stream_closed(&mut stream, check_timeout).await {
                    self.proxy_list.add_address(addr);
                    info!(
                        "Proxy Connect to: {addr} : record [{:.3}s]",
                        elapsed.as_secs_f64()
                    );
                } else {
                    info!(
                        "Proxy Connect to: {addr} : unrecord [{:.3}s]",
                        elapsed.as_secs_f64()
                    );
                }
                Ok(stream)
            }
            Err(e) => Err(e),
        }
    }
}

#[async_trait]
impl Dial for SmartDial {
    async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
        if self.direct_list.contains_address(&addr) {
            debug!("Address {:?} is in direct list, using direct dial", addr);
            match self.direct.dial(addr.clone()).await {
                Ok(stream) => {
                    info!("Direct Connect to: {addr}");
                    return Ok(stream);
                }
                Err(e) => {
                    debug!(
                        "Direct dial failed for {:?}: {}, removing from direct list",
                        addr, e
                    );
                    self.direct_list.remove_address(&addr);
                }
            }
        }

        if self.proxy_list.contains_address(&addr) {
            debug!("Address {:?} is in proxy list, using proxy dial", addr);
            info!("Proxy Connect to: {addr}");
            return self.proxy.dial(addr).await;
        }

        let start = Instant::now();
        let direct_fut = self.direct.dial(addr.clone());
        let proxy_fut = self.proxy.dial(addr.clone());

        tokio::pin!(direct_fut);
        tokio::pin!(proxy_fut);

        tokio::select! {
            direct_res = &mut direct_fut => {
                match direct_res {
                    Ok(stream) => {
                        self.direct_list.add_address(&addr);
                        info!("Direct Connect to: {addr}");
                        Ok(stream)
                    },
                    Err(e) => {
                        debug!("Direct dial failed for {:?}: {}, using proxy", addr, e);
                        let proxy_res = proxy_fut.await;
                        self.handle_proxy_result(proxy_res, &addr, start.elapsed()).await
                    }
                }
            }
            proxy_res = &mut proxy_fut => {
                let direct_res = direct_fut.await;
                match direct_res {
                    Ok(stream) => {
                        self.direct_list.add_address(&addr);
                        info!("Direct Connect to: {addr}");
                        Ok(stream)
                    }
                    Err(e) => {
                        debug!("Direct dial failed for {:?}: {}, using proxy", addr, e);
                        self.handle_proxy_result(proxy_res, &addr, start.elapsed()).await
                    }
                }
            }
        }
    }
}

pub async fn is_stream_closed<R: AsyncRead + Unpin>(
    stream: &mut R,
    timeout_duration: Duration,
) -> bool {
    match timeout(timeout_duration, stream.read(&mut [])).await {
        Ok(Ok(_)) => false,
        Ok(Err(e)) => !matches!(e.kind(), std::io::ErrorKind::WouldBlock),
        Err(_) => false,
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use async_trait::async_trait;
    use socks5_proto::Address;
    use tempfile::NamedTempFile;
    use tokio::{
        io::{AsyncReadExt, AsyncWriteExt},
        net::TcpListener,
        time::Duration,
    };

    use super::{AsyncStream, Dial, DirectDial, SmartDial};
    use crate::address_list::{DirectList, ProxyList};

    struct MockDial {
        succeed: bool,
        delay: Duration,
    }

    impl MockDial {
        fn succeed() -> Self {
            Self {
                succeed: true,
                delay: Duration::ZERO,
            }
        }

        fn fail() -> Self {
            Self {
                succeed: false,
                delay: Duration::ZERO,
            }
        }

        fn with_delay(delay: Duration, succeed: bool) -> Self {
            Self { succeed, delay }
        }

        fn with_timeout_error(delay: Duration) -> Self {
            Self {
                succeed: false,
                delay,
            }
        }
    }

    #[async_trait]
    impl Dial for MockDial {
        async fn dial(&self, _addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
            tokio::time::sleep(self.delay).await;
            if self.succeed {
                Ok(Box::new(tokio::io::duplex(64).0) as Box<dyn AsyncStream>)
            } else {
                Err(anyhow::anyhow!("mock dial failed"))
            }
        }
    }

    #[tokio::test]
    async fn direct_dial_connects_to_local_listener() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let accept_task = tokio::spawn(async move {
            let (mut stream, _) = listener.accept().await.unwrap();
            stream.write_all(b"ok").await.unwrap();
        });

        let mut stream = DirectDial::new(Duration::from_secs(3))
            .dial(Address::SocketAddress(addr))
            .await
            .unwrap();
        let mut buf = [0u8; 2];
        AsyncReadExt::read_exact(&mut stream, &mut buf)
            .await
            .unwrap();

        accept_task.await.unwrap();
        assert_eq!(&buf, b"ok");
    }

    #[tokio::test]
    async fn direct_dial_returns_error_for_unreachable_port() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        drop(listener);

        let result = DirectDial::new(Duration::from_secs(3))
            .dial(Address::SocketAddress(addr))
            .await;

        assert!(result.is_err());
    }

    #[tokio::test]
    async fn smart_dial_uses_proxy_when_domain_in_list() {
        let proxy_temp_file = NamedTempFile::new().unwrap();
        std::fs::write(proxy_temp_file.path(), "blocked.com\n").unwrap();

        let direct_temp_file = NamedTempFile::new().unwrap();

        let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
        let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
        let direct = MockDial::succeed();
        let proxy = MockDial::succeed();

        let smart_dial = SmartDial::new(
            Box::new(direct),
            Box::new(proxy),
            proxy_list,
            direct_list,
            Duration::from_secs(3),
        );

        let result = smart_dial
            .dial(Address::DomainAddress(b"blocked.com".to_vec(), 443))
            .await;

        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn smart_dial_uses_direct_when_succeeds_first() {
        let proxy_temp_file = NamedTempFile::new().unwrap();
        let direct_temp_file = NamedTempFile::new().unwrap();
        let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
        let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));

        let direct = MockDial::with_delay(Duration::from_millis(10), true);
        let proxy = MockDial::with_delay(Duration::from_millis(100), true);

        let smart_dial = SmartDial::new(
            Box::new(direct),
            Box::new(proxy),
            proxy_list,
            direct_list,
            Duration::from_secs(3),
        );

        let result = smart_dial
            .dial(Address::DomainAddress(b"fast-direct.com".to_vec(), 443))
            .await;

        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn smart_dial_adds_domain_when_direct_times_out_and_proxy_succeeds() {
        let proxy_temp_file = NamedTempFile::new().unwrap();
        let direct_temp_file = NamedTempFile::new().unwrap();
        let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
        let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));

        let direct = MockDial::with_timeout_error(Duration::from_millis(100));
        let proxy = MockDial::with_delay(Duration::from_millis(10), true);

        let smart_dial = SmartDial::new(
            Box::new(direct),
            Box::new(proxy),
            proxy_list.clone(),
            direct_list,
            Duration::from_secs(3),
        );

        let result = smart_dial
            .dial(Address::DomainAddress(b"slow-direct.com".to_vec(), 443))
            .await;

        assert!(result.is_ok());
        assert!(
            proxy_list.contains_address(&Address::DomainAddress(b"slow-direct.com".to_vec(), 443))
        );
    }

    #[tokio::test]
    async fn smart_dial_returns_error_when_both_fail() {
        let proxy_temp_file = NamedTempFile::new().unwrap();
        let direct_temp_file = NamedTempFile::new().unwrap();
        let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
        let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));

        let direct = MockDial::fail();
        let proxy = MockDial::fail();

        let smart_dial = SmartDial::new(
            Box::new(direct),
            Box::new(proxy),
            proxy_list,
            direct_list,
            Duration::from_secs(3),
        );

        let result = smart_dial
            .dial(Address::DomainAddress(b"both-fail.com".to_vec(), 443))
            .await;

        assert!(result.is_err());
    }

    #[tokio::test]
    async fn smart_dial_uses_direct_when_domain_in_direct_list() {
        let proxy_temp_file = NamedTempFile::new().unwrap();
        let direct_temp_file = NamedTempFile::new().unwrap();
        std::fs::write(direct_temp_file.path(), "direct.com\n").unwrap();

        let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
        let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
        let direct = MockDial::succeed();
        let proxy = MockDial::fail();

        let smart_dial = SmartDial::new(
            Box::new(direct),
            Box::new(proxy),
            proxy_list,
            direct_list,
            Duration::from_secs(3),
        );

        let result = smart_dial
            .dial(Address::DomainAddress(b"direct.com".to_vec(), 443))
            .await;

        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn smart_dial_removes_from_direct_list_on_failure() {
        let proxy_temp_file = NamedTempFile::new().unwrap();
        let direct_temp_file = NamedTempFile::new().unwrap();
        std::fs::write(direct_temp_file.path(), "direct-fail.com\n").unwrap();

        let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
        let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
        let direct = MockDial::fail();
        let proxy = MockDial::succeed();

        let smart_dial = SmartDial::new(
            Box::new(direct),
            Box::new(proxy),
            proxy_list,
            direct_list.clone(),
            Duration::from_secs(3),
        );

        let result = smart_dial
            .dial(Address::DomainAddress(b"direct-fail.com".to_vec(), 443))
            .await;

        assert!(result.is_ok());
        assert!(
            !direct_list
                .contains_address(&Address::DomainAddress(b"direct-fail.com".to_vec(), 443))
        );
    }

    #[tokio::test]
    async fn smart_dial_adds_to_direct_list_on_success() {
        let proxy_temp_file = NamedTempFile::new().unwrap();
        let direct_temp_file = NamedTempFile::new().unwrap();
        let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
        let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));

        let direct = MockDial::with_delay(Duration::from_millis(10), true);
        let proxy = MockDial::with_delay(Duration::from_millis(100), true);

        let smart_dial = SmartDial::new(
            Box::new(direct),
            Box::new(proxy),
            proxy_list,
            direct_list.clone(),
            Duration::from_secs(3),
        );

        let result = smart_dial
            .dial(Address::DomainAddress(b"new-direct.com".to_vec(), 443))
            .await;

        assert!(result.is_ok());
        assert!(
            direct_list.contains_address(&Address::DomainAddress(b"new-direct.com".to_vec(), 443))
        );
    }
}
