use anyhow::Context;
use log::{debug, error};
use socks5_proto::{
    Address, Command, Reply, Request as SocksRequest, Response,
    handshake::{self, Method},
};
use tokio::{
    io::{AsyncWriteExt, copy_bidirectional},
    net::TcpStream,
};

use crate::dial::Dial;
use crate::proto::http_connect;

pub struct ProxyConnection {
    ts: TcpStream,
    dial: Box<dyn Dial>,
}

impl ProxyConnection {
    pub fn new(ts: TcpStream, dial: Box<dyn Dial>) -> Self {
        Self { ts, dial }
    }

    /// Detect the inbound protocol and proxy the connection until either side closes.
    pub async fn handle(self) {
        let mut first_bit = [0u8];
        if let Err(e) = self.ts.peek(&mut first_bit).await {
            error!("can't peek first_bit err: {e}");
            return;
        }

        let ret = if first_bit[0] == socks5_proto::SOCKS_VERSION {
            self.handle_socks().await
        } else {
            self.handle_http().await
        };
        if let Err(e) = ret {
            error!("proxy handle err: {e:?}");
        };
    }

    async fn handle_socks(mut self) -> anyhow::Result<()> {
        debug!(
            "socks proxy connection {:?} to {:?}",
            self.ts.peer_addr().ok(),
            self.ts.local_addr().ok()
        );

        let _req = handshake::Request::read_from(&mut self.ts)
            .await
            .context("socks handshake failed")?;

        let resp = handshake::Response::new(Method::NONE);

        resp.write_to(&mut self.ts)
            .await
            .context("socks write response failed")?;

        let req = SocksRequest::read_from(&mut self.ts)
            .await
            .context("socks read request failed")?;

        let addr = req.address;

        debug!("start connect {addr}");
        match req.command {
            Command::Connect => {
                let target = self.dial.dial(addr.clone()).await;
                match target {
                    Ok(mut target) => {
                        self.socks_reply(Reply::Succeeded, Address::unspecified())
                            .await?;

                        if let Ok((a, b)) = copy_bidirectional(&mut self.ts, &mut target).await {
                            debug!(
                                "socks copy end for {} traffic: {}<=>{} total: {}",
                                addr,
                                a,
                                b,
                                a + b
                            );
                        }

                        Ok(())
                    }
                    Err(e) => {
                        self.socks_reply(Reply::HostUnreachable, Address::unspecified())
                            .await
                            .context("socks reply failed.")?;
                        Err(e).context(format!("socks dial {addr} failed ."))
                    }
                }
            }
            cmd => {
                debug!("socks unsupported command {:?}", cmd);
                self.socks_reply(Reply::CommandNotSupported, Address::unspecified())
                    .await?;
                Ok(())
            }
        }
    }

    async fn handle_http(mut self) -> anyhow::Result<()> {
        debug!(
            "http proxy connection {:?} to {:?}",
            self.ts.peer_addr().ok(),
            self.ts.local_addr().ok()
        );
        let buf = http_connect::read_http_request_end(&mut self.ts)
            .await
            .context("http proxy read http request end failed")?;

        debug!(
            "http proxy read buf: \n{}",
            String::from_utf8_lossy(buf.as_slice())
        );
        match http_connect::HttpConnectRequest::parse(buf.as_slice()) {
            Ok(req) => {
                let addr = req.addr().clone();
                let mut target = self
                    .dial
                    .dial(addr.clone())
                    .await
                    .context(format!("http proxy connect addr {} failed", addr))?;

                if let Some(data) = req.nugget() {
                    target
                        .write_all(data.data().as_slice())
                        .await
                        .context("http proxy target write_all buf failed")?;
                    target
                        .flush()
                        .await
                        .context("http proxy flush target failed")?;
                } else {
                    self.ts
                        .write("HTTP/1.1 200 OK\r\n\r\n".as_bytes())
                        .await
                        .context("http proxy write response failed")?;
                }

                if let Ok((a, b)) = copy_bidirectional(&mut self.ts, &mut target).await {
                    debug!(
                        "http copy end for {} traffic: {}<=>{} total: {}",
                        addr,
                        a,
                        b,
                        a + b
                    );
                };
                Ok(())
            }
            Err(e) => {
                debug!("http proxy BAD_REQUEST");
                self.ts
                    .write("HTTP/1.1 400 BAD_REQUEST\r\n\r\n".as_bytes())
                    .await
                    .context("http proxy write response failed")?;
                Err(e).context("http dial failed .".to_string())
            }
        }
    }

    async fn socks_reply(&mut self, reply: Reply, addr: Address) -> anyhow::Result<()> {
        let resp = Response::new(reply, addr);
        resp.write_to(&mut self.ts)
            .await
            .context("scoks write reply response failed")
    }
}

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

    use anyhow::anyhow;
    use async_trait::async_trait;
    use socks5_proto::{
        Address, Command, Reply, Request as SocksRequest, Response,
        handshake::{Method, Request as HandshakeRequest, Response as HandshakeResponse},
    };
    use tokio::{
        io::{AsyncReadExt, AsyncWriteExt, DuplexStream},
        net::{TcpListener, TcpStream},
    };

    use super::ProxyConnection;
    use crate::dial::{AsyncStream, Dial};

    struct MockDial {
        result: Mutex<Option<anyhow::Result<DuplexStream>>>,
        seen_addrs: Arc<Mutex<Vec<Address>>>,
    }

    impl MockDial {
        fn succeed(stream: DuplexStream, seen_addrs: Arc<Mutex<Vec<Address>>>) -> Self {
            Self {
                result: Mutex::new(Some(Ok(stream))),
                seen_addrs,
            }
        }

        fn fail(err: anyhow::Error, seen_addrs: Arc<Mutex<Vec<Address>>>) -> Self {
            Self {
                result: Mutex::new(Some(Err(err))),
                seen_addrs,
            }
        }
    }

    #[async_trait]
    impl Dial for MockDial {
        async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
            self.seen_addrs.lock().unwrap().push(addr);
            self.result
                .lock()
                .unwrap()
                .take()
                .expect("dial should only be called once")
                .map(|s| Box::new(s) as Box<dyn AsyncStream>)
        }
    }

    async fn tcp_pair() -> (TcpStream, TcpStream) {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let client = TcpStream::connect(addr).await.unwrap();
        let (server, _) = listener.accept().await.unwrap();

        (server, client)
    }

    #[tokio::test]
    async fn handle_http_connect_replies_ok_and_dials_target() {
        let (server, mut client) = tcp_pair().await;
        let (target, mut target_peer) = tokio::io::duplex(256);
        let seen_addrs = Arc::new(Mutex::new(Vec::new()));
        let dial = MockDial::succeed(target, seen_addrs.clone());

        let proxy = ProxyConnection::new(server, Box::new(dial));
        let proxy_task = tokio::spawn(async move { proxy.handle_http().await });
        let target_task = tokio::spawn(async move {
            let mut buf = Vec::new();
            target_peer.read_to_end(&mut buf).await.unwrap();
            buf
        });

        client
            .write_all(b"CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\n")
            .await
            .unwrap();

        let mut response = [0u8; 19];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(&response, b"HTTP/1.1 200 OK\r\n\r\n");

        client.shutdown().await.unwrap();
        drop(client);

        proxy_task.await.unwrap().unwrap();
        assert!(target_task.await.unwrap().is_empty());
        assert_eq!(
            seen_addrs.lock().unwrap().as_slice(),
            &[Address::DomainAddress(b"example.com".to_vec(), 443)]
        );
    }

    #[tokio::test]
    async fn handle_http_with_nugget_forwards_request_body_to_target() {
        let (server, mut client) = tcp_pair().await;
        let (target, mut target_peer) = tokio::io::duplex(512);
        let seen_addrs = Arc::new(Mutex::new(Vec::new()));
        let dial = MockDial::succeed(target, seen_addrs.clone());
        let raw = b"GET https://upstream.example/path HTTP/1.1\r\nHost: service.internal\r\n\r\n";

        let proxy = ProxyConnection::new(server, Box::new(dial));
        let proxy_task = tokio::spawn(async move { proxy.handle_http().await });
        let target_task = tokio::spawn(async move {
            let mut received = vec![0; raw.len()];
            target_peer.read_exact(&mut received).await.unwrap();
            target_peer
                .write_all(b"HTTP/1.1 204 No Content\r\n\r\n")
                .await
                .unwrap();
            target_peer.shutdown().await.unwrap();
            received
        });

        client.write_all(raw).await.unwrap();

        let mut response = vec![0; 27];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, b"HTTP/1.1 204 No Content\r\n\r\n");

        client.shutdown().await.unwrap();
        drop(client);

        proxy_task.await.unwrap().unwrap();
        assert_eq!(target_task.await.unwrap(), raw);
        assert_eq!(
            seen_addrs.lock().unwrap().as_slice(),
            &[Address::DomainAddress(b"service.internal".to_vec(), 443)]
        );
    }

    #[tokio::test]
    async fn handle_http_bad_request_returns_400_without_dialing() {
        let (server, mut client) = tcp_pair().await;
        let seen_addrs = Arc::new(Mutex::new(Vec::new()));
        let dial = MockDial::fail(anyhow!("dial should not be called"), seen_addrs.clone());

        let proxy = ProxyConnection::new(server, Box::new(dial));
        let proxy_task = tokio::spawn(async move { proxy.handle_http().await });

        client.write_all(b"BAD\r\n\r\n").await.unwrap();

        let mut response = [0u8; 28];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(&response, b"HTTP/1.1 400 BAD_REQUEST\r\n\r\n");

        client.shutdown().await.unwrap();
        drop(client);

        assert!(proxy_task.await.unwrap().is_err());
        assert!(seen_addrs.lock().unwrap().is_empty());
    }

    #[tokio::test]
    async fn handle_socks_connect_negotiates_and_replies_succeeded() {
        let (server, mut client) = tcp_pair().await;
        let (target, mut target_peer) = tokio::io::duplex(256);
        let seen_addrs = Arc::new(Mutex::new(Vec::new()));
        let dial = MockDial::succeed(target, seen_addrs.clone());

        let proxy = ProxyConnection::new(server, Box::new(dial));
        let proxy_task = tokio::spawn(async move { proxy.handle_socks().await });
        let target_task = tokio::spawn(async move {
            let mut buf = Vec::new();
            target_peer.read_to_end(&mut buf).await.unwrap();
            buf
        });

        HandshakeRequest::new(vec![Method::NONE])
            .write_to(&mut client)
            .await
            .unwrap();
        let handshake = HandshakeResponse::read_from(&mut client).await.unwrap();
        assert_eq!(handshake.method, Method::NONE);

        SocksRequest::new(
            Command::Connect,
            Address::DomainAddress(b"example.com".to_vec(), 1080),
        )
        .write_to(&mut client)
        .await
        .unwrap();

        let response = Response::read_from(&mut client).await.unwrap();
        assert_eq!(response.reply, Reply::Succeeded);

        client.shutdown().await.unwrap();
        drop(client);

        proxy_task.await.unwrap().unwrap();
        assert!(target_task.await.unwrap().is_empty());
        assert_eq!(
            seen_addrs.lock().unwrap().as_slice(),
            &[Address::DomainAddress(b"example.com".to_vec(), 1080)]
        );
    }
}
