#[macro_use]
extern crate log;
use futures_util::StreamExt;
use std::future::Future;
use std::io;
use tokio::net::ToSocketAddrs;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;
use tokio::sync::mpsc::Sender;
use tokio_tungstenite::WebSocketStream;
use tungstenite::Message;

pub struct WsEndpoint<T, F> {
    processor: fn(T) -> F,
    sync: bool,
}

impl<F: 'static + Future<Output = Option<String>> + Send> WsEndpoint<String, F> {
    pub fn new_async_text_endpoint(processor: fn(String) -> F) -> WsEndpoint<String, F> {
        WsEndpoint {
            processor,
            sync: false,
        }
    }

    pub fn new_sync_text_endpoint(processor: fn(String) -> F) -> WsEndpoint<String, F> {
        WsEndpoint {
            processor,
            sync: true,
        }
    }

    pub async fn start<A: ToSocketAddrs>(self, addr: A) -> io::Result<()> {
        let mut listener = TcpListener::bind(addr).await?;

        loop {
            info!("Accepting new clients...");
            let incoming = listener.accept().await;
            match incoming {
                Ok((socket, _)) => {
                    tokio::spawn(serve_text_client(socket, self.processor, self.sync));
                }
                Err(e) => {
                    error!("Error accepting client connection: {}", e);
                }
            }
        }
    }
}

impl<F: 'static + Future<Output = Option<Vec<u8>>> + Send> WsEndpoint<Vec<u8>, F> {
    pub fn new_async_binary_endpoint(processor: fn(Vec<u8>) -> F) -> WsEndpoint<Vec<u8>, F> {
        WsEndpoint {
            processor,
            sync: false,
        }
    }

    pub fn new_sync_binary_endpoint(processor: fn(Vec<u8>) -> F) -> WsEndpoint<Vec<u8>, F> {
        WsEndpoint {
            processor,
            sync: true,
        }
    }

    pub async fn start<A: ToSocketAddrs>(self, addr: A) -> io::Result<()> {
        let mut listener = TcpListener::bind(addr).await?;

        loop {
            info!("Accepting new clients...");
            let incoming = listener.accept().await;
            match incoming {
                Ok((socket, _)) => {
                    tokio::spawn(serve_binary_client(socket, self.processor, self.sync));
                }
                Err(e) => {
                    error!("Error accepting client connection: {}", e);
                }
            }
        }
    }
}

async fn serve_text_client<F: 'static + Future<Output = Option<String>> + Send>(
    socket: TcpStream,
    processor: fn(String) -> F,
    sync: bool,
) {
    match tokio_tungstenite::accept_async(socket).await {
        Ok(ws_stream) => {
            process_text_stream(ws_stream, processor, sync).await;
        }
        Err(e) => error!("Error during the websocket handshake occurred: {}", e),
    }
}

async fn serve_binary_client<F: 'static + Future<Output = Option<Vec<u8>>> + Send>(
    socket: TcpStream,
    processor: fn(Vec<u8>) -> F,
    sync: bool,
) {
    match tokio_tungstenite::accept_async(socket).await {
        Ok(ws_stream) => {
            process_binary_stream(ws_stream, processor, sync).await;
        }
        Err(e) => error!("Error during the websocket handshake occurred: {}", e),
    }
}

async fn process_text_stream<F: 'static + Future<Output = Option<String>> + Send>(
    ws_stream: WebSocketStream<TcpStream>,
    processor: fn(String) -> F,
    sync: bool,
) {
    let (write, mut read) = ws_stream.split();

    let (response_tx, response_rx): (Sender<Option<String>>, _) = mpsc::channel(100);

    tokio::spawn(
        response_rx
            .filter_map(|e| async { e.map(|r| Ok(Message::from(r))) })
            .forward(write),
    );

    while let Some(input) = read.next().await {
        match input {
            Ok(input) => {
                if input.is_text() {
                    process_text_request(input, response_tx.clone(), processor, sync).await;
                }
            }
            Err(e) => error!("Error receiving WS message: {}", e),
        }
    }
}

async fn process_binary_stream<F: 'static + Future<Output = Option<Vec<u8>>> + Send>(
    ws_stream: WebSocketStream<TcpStream>,
    processor: fn(Vec<u8>) -> F,
    sync: bool,
) {
    let (write, mut read) = ws_stream.split();

    let (response_tx, response_rx): (Sender<Option<Vec<u8>>>, _) = mpsc::channel(100);

    tokio::spawn(
        response_rx
            .filter_map(|e| async { e.map(|r| Ok(Message::from(r))) })
            .forward(write),
    );

    while let Some(input) = read.next().await {
        match input {
            Ok(input) => {
                dbg!(&input);
                if input.is_binary() {
                    process_binary_request(input, response_tx.clone(), processor, sync).await;
                }
            }
            Err(e) => error!("Error receiving WS message: {}", e),
        }
    }
}

async fn process_text_request<F: 'static + Future<Output = Option<String>> + Send>(
    msg: Message,
    response_tx: Sender<Option<String>>,
    processor: fn(String) -> F,
    sync: bool,
) {
    match msg.into_text() {
        Ok(request) => {
            if sync {
                do_process(request, processor, response_tx).await;
            } else {
                tokio::spawn(do_process(request, processor, response_tx));
            }
        }
        Err(e) => {
            error!("Error converting message to text: {}", e);
            return;
        }
    }
}

async fn process_binary_request<F: 'static + Future<Output = Option<Vec<u8>>> + Send>(
    msg: Message,
    response_tx: Sender<Option<Vec<u8>>>,
    processor: fn(Vec<u8>) -> F,
    sync: bool,
) {
    let request = msg.into_data();
    if sync {
        do_process(request, processor, response_tx).await;
    } else {
        tokio::spawn(do_process(request, processor, response_tx));
    }
}

async fn do_process<T: 'static, F: 'static + Future<Output = Option<T>> + Send>(
    msg: T,
    process: fn(T) -> F,
    response_tx: Sender<Option<T>>,
) {
    let mut response_tx = response_tx;
    let response = process(msg).await;
    if let Err(e) = response_tx.send(response).await {
        error!("Error sending response to WS: {}", e);
    }
}