//! SOCKS5 TCP Server

use std::{
    io::{self, ErrorKind},
    net::{Ipv4Addr, SocketAddr},
    sync::Arc,
};

use log::{debug, error, trace, warn};
use shadowsocks::{
    config::Mode,
    relay::socks5::{
        self,
        Address,
        Command,
        Error as Socks5Error,
        HandshakeRequest,
        HandshakeResponse,
        Reply,
        TcpRequestHeader,
        TcpResponseHeader,
    },
    ServerAddr,
};
use tokio::net::TcpStream;

use crate::{
    local::{
        context::ServiceContext,
        loadbalancing::PingBalancer,
        net::AutoProxyClientStream,
        utils::establish_tcp_tunnel,
    },
    net::utils::ignore_until_end,
};

pub struct Socks5TcpHandler {
    context: Arc<ServiceContext>,
    udp_bind_addr: Option<Arc<ServerAddr>>,
    balancer: PingBalancer,
    mode: Mode,
}

impl Socks5TcpHandler {
    pub fn new(
        context: Arc<ServiceContext>,
        udp_bind_addr: Option<Arc<ServerAddr>>,
        balancer: PingBalancer,
        mode: Mode,
    ) -> Socks5TcpHandler {
        Socks5TcpHandler {
            context,
            udp_bind_addr,
            balancer,
            mode,
        }
    }

    pub async fn handle_socks5_client(self, mut stream: TcpStream, peer_addr: SocketAddr) -> io::Result<()> {
        // 1. Handshake

        let handshake_req = match HandshakeRequest::read_from(&mut stream).await {
            Ok(r) => r,
            Err(Socks5Error::IoError(ref err)) if err.kind() == ErrorKind::UnexpectedEof => {
                trace!("socks5 handshake early eof. peer: {}", peer_addr);
                return Ok(());
            }
            Err(err) => {
                error!("socks5 handshake error: {}", err);
                return Err(err.into());
            }
        };

        trace!("socks5 {:?}", handshake_req);

        if !handshake_req.methods.contains(&socks5::SOCKS5_AUTH_METHOD_NONE) {
            use std::io::Error;

            let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE);
            resp.write_to(&mut stream).await?;

            return Err(Error::new(
                ErrorKind::Other,
                "currently shadowsocks-rust does not support authentication",
            ));
        } else {
            // Reply to client
            let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NONE);
            trace!("reply handshake {:?}", resp);
            resp.write_to(&mut stream).await?;
        }

        // 2. Fetch headers
        let header = match TcpRequestHeader::read_from(&mut stream).await {
            Ok(h) => h,
            Err(err) => {
                error!("failed to get TcpRequestHeader: {}, peer: {}", err, peer_addr);
                let rh = TcpResponseHeader::new(err.as_reply(), Address::SocketAddress(peer_addr));
                rh.write_to(&mut stream).await?;
                return Err(err.into());
            }
        };

        trace!("socks5 {:?} peer: {}", header, peer_addr);

        let addr = header.address;

        // 3. Handle Command
        match header.command {
            Command::TcpConnect => {
                debug!("CONNECT {}", addr);

                self.handle_tcp_connect(stream, peer_addr, addr).await
            }
            Command::UdpAssociate => {
                debug!("UDP ASSOCIATE from {}", addr);

                self.handle_udp_associate(stream, addr).await
            }
            Command::TcpBind => {
                warn!("BIND is not supported");
                let rh = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr);
                rh.write_to(&mut stream).await?;

                Ok(())
            }
        }
    }

    async fn handle_tcp_connect(
        self,
        mut stream: TcpStream,
        peer_addr: SocketAddr,
        target_addr: Address,
    ) -> io::Result<()> {
        if !self.mode.enable_tcp() {
            warn!("TCP CONNECT is disabled");

            let rh = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, target_addr);
            rh.write_to(&mut stream).await?;

            return Ok(());
        }

        let server = self.balancer.best_tcp_server();
        let svr_cfg = server.server_config();

        let mut remote = match AutoProxyClientStream::connect(self.context.clone(), &server, &target_addr).await {
            Ok(remote) => {
                // Tell the client that we are ready
                let header =
                    TcpResponseHeader::new(socks5::Reply::Succeeded, Address::SocketAddress(remote.local_addr()?));
                header.write_to(&mut stream).await?;

                trace!("sent header: {:?}", header);

                remote
            }
            Err(err) => {
                let reply = match err.kind() {
                    ErrorKind::ConnectionRefused => Reply::ConnectionRefused,
                    ErrorKind::ConnectionAborted => Reply::HostUnreachable,
                    _ => Reply::NetworkUnreachable,
                };

                let dummy_address = SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0);
                let header = TcpResponseHeader::new(reply, Address::SocketAddress(dummy_address));
                header.write_to(&mut stream).await?;

                return Err(err);
            }
        };

        establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, &target_addr).await
    }

    async fn handle_udp_associate(self, mut stream: TcpStream, client_addr: Address) -> io::Result<()> {
        match self.udp_bind_addr {
            None => {
                warn!("socks5 udp is disabled");

                let rh = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, client_addr);
                rh.write_to(&mut stream).await?;

                Ok(())
            }
            Some(bind_addr) => {
                // shadowsocks accepts both TCP and UDP from the same address

                let rh = TcpResponseHeader::new(socks5::Reply::Succeeded, bind_addr.as_ref().into());
                rh.write_to(&mut stream).await?;

                // Hold connection until EOF.
                let _ = ignore_until_end(&mut stream).await;

                Ok(())
            }
        }
    }
}
