use std::{
    collections::HashMap,
    pin::Pin,
    task::{Context, Poll},
};

use anyhow::Result;
use async_tungstenite::{tungstenite::protocol::Role, WebSocketStream};
use base64::{engine::general_purpose::STANDARD as STANDARD_BASE64, Engine};
use futures::{AsyncRead, AsyncWrite, SinkExt, StreamExt};
use libcfd::{
    connection::{
        ConnectRequest, ConnectResponse, Connection, DstAddr, QuinnRecvStream, QuinnSendStream,
        SrcAddr,
    },
    tunnel_config::TunnelConfig,
};
use sha1::{Digest, Sha1};
use tokio::{
    runtime::Builder as RuntimeBuilder,
    task::{spawn_local, LocalSet},
};

fn main() {
    let runtime = RuntimeBuilder::new_current_thread()
        .enable_all()
        .build()
        .expect("failed to build tokio runtime");
    runtime.block_on(LocalSet::new().run_until(async {
        if let Err(error) = async_main().await {
            println!("{error}");
        }
    }));
}

async fn async_main() -> Result<()> {
    let config = TunnelConfig::try_cloudflare().await?;
    println!("Quick tunnel created at wss://{}", config.hostname);
    let (connection, connection_details) =
        Connection::new(config, 0, SrcAddr::Default, DstAddr::Default).await?;
    println!("Registered tunnel connection: {:#?}", connection_details);
    loop {
        let connect_request = connection.accept().await?;
        spawn_local(async move {
            if let Err(error) = handle_connect_request(connect_request).await {
                println!("{error}")
            }
        });
    }
}

async fn handle_connect_request(connect_request: ConnectRequest) -> Result<()> {
    println!("Incoming request: {:#?}", connect_request.metadata);
    let mut metadata = HashMap::new();
    metadata.insert("HttpStatus".to_string(), "101".to_string());
    let mut digest = Sha1::new();
    digest.update(&connect_request.metadata["HttpHeader:Sec-Websocket-Key"]);
    digest.update("258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
    let accept = STANDARD_BASE64.encode(digest.finalize());
    metadata.insert("HttpHeader:Sec-Websocket-Accept".to_string(), accept);
    let (send_stream, recv_stream) = connect_request
        .respond_with(ConnectResponse::Metadata(metadata))
        .await?;
    let mut websocket = WebSocketStream::from_raw_socket(
        Stream {
            send_stream,
            recv_stream,
        },
        Role::Server,
        None,
    )
    .await;
    while let Some(Ok(message)) = websocket.next().await {
        websocket.send(message).await?
    }
    Ok(())
}

struct Stream {
    send_stream: QuinnSendStream,
    recv_stream: QuinnRecvStream,
}

impl AsyncWrite for Stream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        AsyncWrite::poll_write(Box::pin(&mut self.send_stream).as_mut(), cx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        AsyncWrite::poll_flush(Box::pin(&mut self.send_stream).as_mut(), cx)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        AsyncWrite::poll_close(Box::pin(&mut self.send_stream).as_mut(), cx)
    }
}

impl AsyncRead for Stream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<std::io::Result<usize>> {
        AsyncRead::poll_read(Box::pin(&mut self.recv_stream).as_mut(), cx, buf)
    }
}