use std::fmt;
use std::io;

use base64::engine::general_purpose::STANDARD;
use base64::Engine;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use url::Url;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ProxyKind {
    Http,
    Socks5,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ProxyConfig {
    pub kind: ProxyKind,
    pub host: String,
    pub port: u16,
    pub username: Option<String>,
    pub password: Option<String>,
}

#[derive(Debug)]
pub enum ProxyError {
    Invalid(String),
    Io(io::Error),
    Rejected(String),
}

impl fmt::Display for ProxyError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Invalid(message) => write!(formatter, "invalid proxy: {message}"),
            Self::Io(error) => error.fmt(formatter),
            Self::Rejected(message) => write!(formatter, "proxy rejected connection: {message}"),
        }
    }
}

impl std::error::Error for ProxyError {}

impl From<io::Error> for ProxyError {
    fn from(error: io::Error) -> Self {
        Self::Io(error)
    }
}

impl ProxyConfig {
    pub fn parse(source: &str) -> Result<Self, ProxyError> {
        let url = Url::parse(source).map_err(|error| ProxyError::Invalid(error.to_string()))?;
        let kind = match url.scheme() {
            "http" => ProxyKind::Http,
            "socks5" | "socks5h" => ProxyKind::Socks5,
            scheme => return Err(ProxyError::Invalid(format!("unsupported scheme {scheme}"))),
        };
        let host = url
            .host_str()
            .filter(|host| !host.is_empty())
            .ok_or_else(|| ProxyError::Invalid("host is missing".to_string()))?
            .to_string();
        let port = url
            .port_or_known_default()
            .ok_or_else(|| ProxyError::Invalid("port is missing".to_string()))?;
        let username = (!url.username().is_empty()).then(|| url.username().to_string());
        let password = url.password().map(str::to_string);
        Ok(Self {
            kind,
            host,
            port,
            username,
            password,
        })
    }
}

pub async fn connect(
    target_host: &str,
    target_port: u16,
    proxy: Option<&ProxyConfig>,
) -> Result<TcpStream, ProxyError> {
    let Some(proxy) = proxy else {
        return Ok(TcpStream::connect((target_host, target_port)).await?);
    };
    let mut stream = TcpStream::connect((proxy.host.as_str(), proxy.port)).await?;
    match proxy.kind {
        ProxyKind::Http => connect_http(&mut stream, target_host, target_port, proxy).await?,
        ProxyKind::Socks5 => connect_socks5(&mut stream, target_host, target_port, proxy).await?,
    }
    Ok(stream)
}

async fn connect_http(
    stream: &mut TcpStream,
    host: &str,
    port: u16,
    proxy: &ProxyConfig,
) -> Result<(), ProxyError> {
    let mut request = format!("CONNECT {host}:{port} HTTP/1.1\r\nHost: {host}:{port}\r\n");
    if let Some(username) = &proxy.username {
        let credentials = format!(
            "{username}:{}",
            proxy.password.as_deref().unwrap_or_default()
        );
        request.push_str(&format!(
            "Proxy-Authorization: Basic {}\r\n",
            STANDARD.encode(credentials)
        ));
    }
    request.push_str("Connection: keep-alive\r\n\r\n");
    stream.write_all(request.as_bytes()).await?;
    stream.flush().await?;
    let mut response = Vec::with_capacity(512);
    let mut byte = [0u8; 1];
    while !response.ends_with(b"\r\n\r\n") {
        if response.len() >= 16 * 1024 {
            return Err(ProxyError::Rejected(
                "response header is too large".to_string(),
            ));
        }
        stream.read_exact(&mut byte).await?;
        response.push(byte[0]);
    }
    let first_line = String::from_utf8_lossy(&response)
        .lines()
        .next()
        .unwrap_or_default()
        .to_string();
    if first_line.split_whitespace().nth(1) != Some("200") {
        return Err(ProxyError::Rejected(first_line));
    }
    Ok(())
}

async fn connect_socks5(
    stream: &mut TcpStream,
    host: &str,
    port: u16,
    proxy: &ProxyConfig,
) -> Result<(), ProxyError> {
    let methods: &[u8] = if proxy.username.is_some() {
        &[0, 2]
    } else {
        &[0]
    };
    let mut greeting = vec![5, methods.len() as u8];
    greeting.extend_from_slice(methods);
    stream.write_all(&greeting).await?;
    let mut choice = [0u8; 2];
    stream.read_exact(&mut choice).await?;
    if choice[0] != 5 || choice[1] == 0xff {
        return Err(ProxyError::Rejected("authentication method".to_string()));
    }
    if choice[1] == 2 {
        authenticate_socks5(stream, proxy).await?;
    } else if choice[1] != 0 {
        return Err(ProxyError::Rejected(
            "unsupported authentication method".to_string(),
        ));
    }
    let host_bytes = host.as_bytes();
    if host_bytes.len() > u8::MAX as usize {
        return Err(ProxyError::Invalid(
            "target hostname is too long".to_string(),
        ));
    }
    let mut request = vec![5, 1, 0, 3, host_bytes.len() as u8];
    request.extend_from_slice(host_bytes);
    request.extend_from_slice(&port.to_be_bytes());
    stream.write_all(&request).await?;
    let mut header = [0u8; 4];
    stream.read_exact(&mut header).await?;
    if header[0] != 5 || header[1] != 0 {
        return Err(ProxyError::Rejected(format!("SOCKS5 status {}", header[1])));
    }
    let address_size = match header[3] {
        1 => 4,
        4 => 16,
        3 => {
            let mut length = [0u8; 1];
            stream.read_exact(&mut length).await?;
            length[0] as usize
        }
        kind => return Err(ProxyError::Rejected(format!("address type {kind}"))),
    };
    let mut remainder = vec![0u8; address_size + 2];
    stream.read_exact(&mut remainder).await?;
    Ok(())
}

async fn authenticate_socks5(
    stream: &mut TcpStream,
    proxy: &ProxyConfig,
) -> Result<(), ProxyError> {
    let username = proxy.username.as_deref().unwrap_or_default().as_bytes();
    let password = proxy.password.as_deref().unwrap_or_default().as_bytes();
    if username.len() > u8::MAX as usize || password.len() > u8::MAX as usize {
        return Err(ProxyError::Invalid("credentials are too long".to_string()));
    }
    let mut request = vec![1, username.len() as u8];
    request.extend_from_slice(username);
    request.push(password.len() as u8);
    request.extend_from_slice(password);
    stream.write_all(&request).await?;
    let mut response = [0u8; 2];
    stream.read_exact(&mut response).await?;
    if response != [1, 0] {
        return Err(ProxyError::Rejected("authentication failed".to_string()));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{ProxyConfig, ProxyKind};

    #[test]
    fn parses_http_credentials() {
        let proxy = ProxyConfig::parse("http://test-user:test-pass@127.0.0.1:8080").unwrap();
        assert_eq!(proxy.kind, ProxyKind::Http);
        assert_eq!(proxy.host, "127.0.0.1");
        assert_eq!(proxy.port, 8080);
        assert_eq!(proxy.username.as_deref(), Some("test-user"));
    }

    #[test]
    fn rejects_unsupported_scheme() {
        assert!(ProxyConfig::parse("ftp://127.0.0.1:21").is_err());
    }
}
