use futures::future::Either;
use reqwest::{
    self,
    header::{HeaderValue, HOST, USER_AGENT},
    Client, Request, Url,
};
use socks5_impl::client;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};

pub struct IncrementalStats {
    count: u64,
    mean: f64,
    m2: f64,
}

impl IncrementalStats {
    pub fn new() -> Self {
        IncrementalStats {
            count: 0,
            mean: 0.0,
            m2: 0.0,
        }
    }

    pub fn add(&mut self, value: f64) {
        self.count += 1;
        let delta = value - self.mean;
        self.mean += delta / self.count as f64;
        let delta2 = value - self.mean;
        self.m2 += delta * delta2;
    }

    pub fn average(&self) -> f64 {
        self.mean
    }

    pub fn variance(&self) -> f64 {
        if self.count < 2 {
            0.0
        } else {
            self.m2 / self.count as f64
        }
    }
}

fn create_http_req(client: reqwest::Client, end_point: &str) -> Option<Request> {
    let url: Url = end_point.parse().ok()?;
    let host = url.host().unwrap().to_string();
    let mut req = client
        .post(url.clone())
        .timeout(Duration::from_secs(5))
        .build()
        .ok()?;
    req.headers_mut()
        .insert(HOST, HeaderValue::from_str(&host).ok()?);
    req.headers_mut()
        .insert(USER_AGENT, HeaderValue::from_static("curl/8.9.1-DEV"));

    Some(req)
}

fn request_to_http_text(req: &Request) -> anyhow::Result<String> {
    use std::fmt::Write;
    let mut result = String::new();
    // Write the request line
    write!(
        &mut result,
        "{} {} {:?}\r\n",
        req.method(),
        req.url(),
        req.version()
    )?;

    // Write the headers
    for (key, value) in req.headers() {
        write!(&mut result, "{}: {}\r\n", key, value.to_str()?)?;
    }

    // End of headers
    write!(&mut result, "\r\n")?;
    if let Some(body) = req.body() {
        write!(
            &mut result,
            "{}",
            String::from_utf8(body.as_bytes().unwrap().to_owned())?
        )?;
    }

    Ok(result)
}

async fn connect_to(url: &str, proxy: Option<&str>) -> anyhow::Result<tokio::net::TcpStream> {
    let url: Url = url.parse()?;
    if let Some(proxy) = proxy {
        let proxy_url: Url = proxy.parse()?;
        let mut ac = tokio::net::TcpStream::connect((
            proxy_url.host().unwrap().to_string(),
            proxy_url.port().unwrap(),
        ))
        .await?;
        client::connect(&mut ac, (url.host().unwrap().to_string(), 443), None).await?;
        Ok(ac)
    } else {
        let port = url.port_or_known_default().unwrap();
        let ac = tokio::net::TcpStream::connect((url.host().unwrap().to_string(), port)).await?;
        Ok(ac)
    }
}
async fn rustls_req(url_str: &str, proxy: Option<&str>) -> anyhow::Result<std::time::Duration> {
    use rustls_pki_types::ServerName;
    use std::sync::Arc;
    use tokio_rustls::rustls::{ClientConfig, RootCertStore};
    use tokio_rustls::TlsConnector;

    // ...

    let mut root_cert_store = RootCertStore::empty();
    root_cert_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
    let config = ClientConfig::builder()
        .with_root_certificates(root_cert_store)
        .with_no_client_auth();
    let connector = TlsConnector::from(Arc::new(config));
    let url: Url = url_str.parse()?;
    let dnsname = ServerName::try_from(url.host().unwrap().to_string())?;

    // let stream = TcpStream::connect(&addr).await?;
    // let mut stream = connector.connect(dnsname, stream).await?;
    // let proxy_url: Url = proxy.parse()?;
    let now = std::time::Instant::now();
    // let mut ac = tokio::net::TcpStream::connect((
    //     proxy_url.host().unwrap().to_string(),
    //     proxy_url.port().unwrap(),
    // ))
    // .await?;
    // log::debug!("tcp conn established time used: {:?}", now.elapsed());
    // let r = client::connect(&mut ac, (url.host().unwrap().to_string(), 443), None).await?;
    // log::debug!("socks5 conn established time used: {:?}", now.elapsed());
    let ac = connect_to(url_str, proxy).await?;
    // let tls_conn = TlsConnector::from(native_tls::TlsConnector::new().unwrap());
    let mut tls_s = connector.connect(dnsname, ac).await?;
    // let url: Url = "https://cp.cloudflare.com/generate_204".parse().unwrap();
    // let mut tls_s = tls_conn.connect("cp.cloudflare.com", ac).await?;
    log::debug!("tls conn established time used: {:?}", now.elapsed());
    let req = create_http_req(Client::new(), url_str).unwrap();
    let txt = request_to_http_text(&req).unwrap();
    tls_s.write_all(&txt.as_bytes()).await?;
    tls_s.flush().await?;
    log::debug!("req write time used: {:?}", now.elapsed());
    let mut buf = vec![0; 4096];
    let n = tls_s.read(&mut buf).await?;
    log::debug!("resp get time used: {:?}, resp len: {}", now.elapsed(), n);
    Ok(now.elapsed())
}

async fn http_req(url_str: &str, proxy: Option<&str>) -> anyhow::Result<std::time::Duration> {
    let url: Url = url_str.parse()?;

    // let stream = TcpStream::connect(&addr).await?;
    // let mut stream = connector.connect(dnsname, stream).await?;
    // let proxy_url: Url = proxy.parse()?;
    let now = std::time::Instant::now();
    // let mut ac = tokio::net::TcpStream::connect((
    //     proxy_url.host().unwrap().to_string(),
    //     proxy_url.port().unwrap(),
    // ))
    // .await?;
    let mut ac = connect_to(url_str, proxy).await?;
    log::debug!("tcp conn established time used: {:?}", now.elapsed());
    let r = client::connect(
        &mut ac,
        (url.host().unwrap().to_string(), url.port().unwrap()),
        None,
    )
    .await?;
    log::debug!("socks5 conn established time used: {:?}", now.elapsed());
    // let tls_conn = TlsConnector::from(native_tls::TlsConnector::new().unwrap());
    // let url: Url = "https://cp.cloudflare.com/generate_204".parse().unwrap();
    // let mut tls_s = tls_conn.connect("cp.cloudflare.com", ac).await?;
    log::debug!("tls conn established time used: {:?}", now.elapsed());
    let req = create_http_req(Client::new(), url_str).unwrap();
    let txt = request_to_http_text(&req).unwrap();
    ac.write_all(&txt.as_bytes()).await?;
    ac.flush().await?;
    log::debug!("req write time used: {:?}", now.elapsed());
    let mut buf = vec![0; 4096];
    let n = ac.read(&mut buf).await?;
    log::debug!("resp get time used: {:?}, resp len: {}", now.elapsed(), n);
    Ok(now.elapsed())
}
#[tokio::main]
async fn main() {
    let url = std::env::args()
        .nth(1)
        .unwrap_or("https://www.youtube.com".to_owned()); // 你可以替换为你想要测量的网址
                                                          // let mut time_rec = Vec::new();
    let iter_num: usize = std::env::args()
        .nth(2)
        .unwrap_or_else(|| "20".to_owned())
        .parse()
        .expect("second parameter should interger");
    let proxy = std::env::args().nth(3);
    let fmt = std::env::args().nth(4);

    let mut stats = IncrementalStats::new();
    let mut succ = 0;
    let mut err = 0;
    for i in 0..iter_num {
        let fut = if url.starts_with("https") {
            Either::Left(rustls_req(&url, proxy.as_deref()))
        } else {
            Either::Right(http_req(&url, proxy.as_deref()))
        };
        match tokio::time::timeout(std::time::Duration::from_secs(5), fut).await {
            Ok(Ok(duration)) => {
                // if response.status().is_success() {
                // let mut duration = start.elapsed();
                // let _r = response.bytes().unwrap();
                if let Some(_name) = fmt.as_ref() {
                    // println!("{name}, {i},{}", duration.as_millis())
                    // duration /= 2;
                } else {
                    println!(
                        "第{i}次测试, 访问 {} 花费了 {:?} 毫秒",
                        url,
                        duration.as_millis()
                    );
                }
                succ += 1;
                // time_rec.push(duration);
                stats.add(duration.as_millis() as f64);
                // }
            }
            Ok(Err(error)) => {
                err += 1;

                if let Some(_name) = fmt.as_ref() {
                    return;
                } else {
                    println!("第{i}次测试, 请求错误：{}", error);
                }
            }
            Err(error) => {
                err += 1;

                if let Some(_name) = fmt.as_ref() {
                    return;
                } else {
                    println!("第{i}次测试, 请求错误：{}", error);
                }
            }
        }
    }

    if let Some(name) = fmt.as_ref() {
        println!(
            "{name}, {},{},{},{}",
            stats.average(),
            stats.variance(),
            succ,
            err
        );
    } else {
        println!(
            "平均: {:?}ms, 方差: {:?} 成功:{} 失败: {}",
            stats.average(),
            stats.variance(),
            succ,
            err
        );
    }
}