use anyhow::{Context, bail};
use bytes::Bytes;
use http_body_util::{BodyExt, Full, combinators::BoxBody};
use hyper::body::Incoming;
use hyper::header::HeaderValue;
use hyper::service::service_fn;
use hyper::{Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use rcgen::{BasicConstraints, CertificateParams, DnType, ExtendedKeyUsagePurpose, IsCa, Issuer, KeyPair, KeyUsagePurpose};
use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
use std::collections::HashMap;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::{Arc, RwLock};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::task::JoinHandle;
use tokio_rustls::TlsAcceptor;

type Body = BoxBody<Bytes, hyper::Error>;

pub struct Task(pub JoinHandle<()>);

impl Drop for Task {
    fn drop(&mut self) {
        self.0.abort();
    }
}

#[derive(Clone, Default)]
pub struct Routes(Arc<RwLock<HashMap<String, SocketAddr>>>);

impl Routes {
    pub fn set(&self, host: &str, backend: SocketAddr) {
        self.0.write().unwrap().insert(host.to_ascii_lowercase(), backend);
    }

    fn get(&self, host: &str) -> Option<SocketAddr> {
        let host = host.to_ascii_lowercase();
        let routes = self.0.read().unwrap();
        routes
            .get(&host)
            .or_else(|| host.split_once('.').and_then(|(_, parent)| routes.get(&format!("*.{parent}"))))
            .copied()
    }
}

pub struct Certs {
    pub ca_pem: String,
    leaf_der: Vec<u8>,
    leaf_key: Vec<u8>,
}

impl Certs {
    pub fn generate(hosts: &[&str]) -> anyhow::Result<Certs> {
        let ca_key = KeyPair::generate()?;
        let mut ca_params = CertificateParams::default();
        ca_params.distinguished_name.push(DnType::CommonName, "zkwid devnet CA");
        ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
        ca_params.key_usages = vec![KeyUsagePurpose::KeyCertSign, KeyUsagePurpose::CrlSign];
        let ca_cert = ca_params.self_signed(&ca_key)?;
        let issuer = Issuer::new(ca_params, ca_key);
        let leaf_key = KeyPair::generate()?;
        let mut leaf_params = CertificateParams::new(hosts.iter().map(|host| host.to_string()).collect::<Vec<_>>())?;
        leaf_params.distinguished_name.push(DnType::CommonName, hosts[0]);
        leaf_params.key_usages = vec![KeyUsagePurpose::DigitalSignature];
        leaf_params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth];
        let leaf = leaf_params.signed_by(&leaf_key, &issuer)?;
        Ok(Certs { ca_pem: ca_cert.pem(), leaf_der: leaf.der().to_vec(), leaf_key: leaf_key.serialize_der() })
    }

    pub fn acceptor(&self) -> anyhow::Result<TlsAcceptor> {
        let config = rustls::ServerConfig::builder_with_provider(Arc::new(rustls::crypto::aws_lc_rs::default_provider()))
            .with_safe_default_protocol_versions()?
            .with_no_client_auth()
            .with_single_cert(
                vec![CertificateDer::from(self.leaf_der.clone())],
                PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(self.leaf_key.clone())),
            )?;
        Ok(TlsAcceptor::from(Arc::new(config)))
    }
}

pub async fn start(routes: Routes, tls: TlsAcceptor) -> anyhow::Result<(SocketAddr, Task)> {
    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let addr = listener.local_addr()?;
    let task = tokio::spawn(async move {
        loop {
            let Ok((stream, peer)) = listener.accept().await else { continue };
            let (routes, tls) = (routes.clone(), tls.clone());
            tokio::spawn(async move {
                tunnel(stream, peer, routes, tls).await.ok();
            });
        }
    });
    Ok((addr, Task(task)))
}

async fn tunnel(mut stream: TcpStream, peer: SocketAddr, routes: Routes, tls: TlsAcceptor) -> anyhow::Result<()> {
    let head = read_head(&mut stream).await?;
    let target = head.lines().next().and_then(|line| line.strip_prefix("CONNECT ")).and_then(|rest| rest.split(' ').next());
    let host = match target.and_then(|target| target.strip_suffix(":443")) {
        Some(host) if routes.get(host).is_some() => host.to_ascii_lowercase(),
        _ => {
            stream.write_all(b"HTTP/1.1 502 Bad Gateway\r\ncontent-length: 0\r\nconnection: close\r\n\r\n").await?;
            return Ok(());
        }
    };
    stream.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n").await?;
    let stream = tls.accept(stream).await?;
    let service = service_fn(move |request| forward(routes.clone(), host.clone(), peer, request));
    hyper::server::conn::http1::Builder::new().serve_connection(TokioIo::new(stream), service).await?;
    Ok(())
}

async fn read_head(stream: &mut TcpStream) -> anyhow::Result<String> {
    let mut head = Vec::new();
    while !head.ends_with(b"\r\n\r\n") {
        if head.len() > 8192 {
            bail!("CONNECT head too large");
        }
        head.push(stream.read_u8().await.context("reading CONNECT head")?);
    }
    Ok(String::from_utf8(head)?)
}

async fn forward(routes: Routes, host: String, peer: SocketAddr, request: Request<Incoming>) -> Result<Response<Body>, Infallible> {
    let Some(backend) = routes.get(&host) else {
        return Ok(bad_gateway(format!("no route for {host}")));
    };
    Ok(match send(backend, &host, peer, request).await {
        Ok(response) => response.map(BodyExt::boxed),
        Err(error) => bad_gateway(format!("{host} -> {backend}: {error:#}")),
    })
}

async fn send(backend: SocketAddr, host: &str, peer: SocketAddr, mut request: Request<Incoming>) -> anyhow::Result<Response<Incoming>> {
    let headers = request.headers_mut();
    headers.insert("x-forwarded-proto", HeaderValue::from_static("https"));
    headers.insert("x-forwarded-host", HeaderValue::from_str(host)?);
    headers.insert("x-forwarded-for", HeaderValue::from_str(&peer.ip().to_string())?);
    let stream = TcpStream::connect(backend).await?;
    let (mut sender, connection) = hyper::client::conn::http1::handshake(TokioIo::new(stream)).await?;
    tokio::spawn(connection);
    Ok(sender.send_request(request).await?)
}

fn bad_gateway(message: String) -> Response<Body> {
    let mut response = Response::new(Full::new(Bytes::from(message)).map_err(|never| match never {}).boxed());
    *response.status_mut() = StatusCode::BAD_GATEWAY;
    response
}