use core::net::SocketAddr;
use core::time::Duration;
use std::path::PathBuf;
use std::process::Stdio;
use std::{env, error};

use axum::body::Body;
use axum::http::{HeaderValue, Request, header};
use axum::middleware::Next;
use axum::response::Response;
use axum::{Router, middleware};
use hyper_util::rt::{TokioExecutor, TokioIo};
use hyper_util::server::conn::auto;
use hyper_util::service::TowerToHyperService;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::process::Command;
use tokio::sync::OnceCell;
use tokio::sync::mpsc::{self, Receiver};
use tokio::task::JoinSet;
use tower_http::services::ServeDir;

const MAX_CONCURRENT: usize = 100;
const MAX_BUFFERED: usize = 500;
const DEFAULT_BIN_PATH: &str = "../build/target.exe";
const DEFAULT_STATIC_PATH: &str = "../frontend";
const TIMEOUT_SECS: u64 = 10;

static BIN_PATH: OnceCell<PathBuf> = OnceCell::const_new();
static STATIC_PATH: OnceCell<PathBuf> = OnceCell::const_new();

#[derive(Debug)]
#[allow(dead_code)]
struct NetworkConnection(pub TcpStream, pub SocketAddr);

#[tokio::main]
async fn main() -> Result<(), Box<dyn error::Error>> {
    if let Ok(path) = env::var("BIN_PATH")
        && path != ""
    {
        BIN_PATH.set(PathBuf::from(path)).unwrap();
    } else {
        BIN_PATH.set(PathBuf::from(DEFAULT_BIN_PATH)).unwrap();
    }

    if let Ok(path) = env::var("STATIC_PATH")
        && path != ""
    {
        STATIC_PATH.set(PathBuf::from(path)).unwrap();
    } else {
        STATIC_PATH.set(PathBuf::from(DEFAULT_STATIC_PATH)).unwrap();
    }

    Command::new(BIN_PATH.get().unwrap())
        .arg("-роанапур=истина")
        .spawn()?
        .wait()
        .await?;

    let listener = TcpListener::bind("0.0.0.0:1337").await?;
    let (tx, rx) = mpsc::channel::<NetworkConnection>(MAX_BUFFERED);

    tokio::spawn(process_connections(rx));

    loop {
        let connection = listener.accept().await?;
        if tx
            .try_send(NetworkConnection(connection.0, connection.1))
            .is_err()
        {
            println!(
                "Connection queue full, dropping connection from {}",
                connection.1
            );
        }
    }
}

async fn process_connections(mut rx: Receiver<NetworkConnection>) {
    let mut join_set = JoinSet::new();

    loop {
        let available = MAX_CONCURRENT - join_set.len();
        tokio::select! {
            connection = rx.recv(), if available > 0 => {
                if let Some(connection) = connection {
                    join_set.spawn(handle_connection_or_timeout(connection));
                }
            },
            _ = join_set.join_next(), if join_set.len() > 0 => {}
        }
    }
}

async fn handle_connection_or_timeout(
    connection: NetworkConnection,
) -> Result<(), Box<dyn error::Error + Send + Sync>> {
    tokio::select! {
        result = handle_connection(connection) => {
            result
        },
        _ = tokio::time::sleep(Duration::from_secs(TIMEOUT_SECS)) => {
            Err("Connection timed out".into())
        }
    }
}

async fn handle_connection(
    mut connection: NetworkConnection,
) -> Result<(), Box<dyn error::Error + Send + Sync>> {
    let is_frontend = check_is_frontend_request(&mut connection).await?;

    if is_frontend {
        return serve_frontend_file(connection).await;
    }

    let mut child = Command::new(BIN_PATH.get().unwrap())
        .arg("-подшефный=истина")
        .stdin(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()?;

    let mut stderr = child.stderr.take().unwrap();
    let mut stdin = child.stdin.take().unwrap();

    loop {
        let mut tcp_buf = [0u8; 4096];
        let mut stderr_buf = [0u8; 4096];

        tokio::select! {
            result = connection.0.read(&mut tcp_buf) => match result {
                Ok(how_many) => stdin.write_all(&tcp_buf[0..how_many]).await?,
                Err(e) => break Err(e.into()),
            },
            result = stderr.read(&mut stderr_buf) => match result {
                Ok(0) => break Ok(()),
                Ok(how_many) => connection.0.write_all(&stderr_buf[0..how_many]).await?,
                Err(e) => break Err(e.into()),
            },
        }
    }
}

async fn check_is_frontend_request(
    connection: &mut NetworkConnection,
) -> Result<bool, Box<dyn error::Error + Send + Sync>> {
    let mut peek_buf = [0u8; 8];
    let bytes_read = connection.0.peek(&mut peek_buf).await?;

    if bytes_read == 0 {
        return Err("Connection closed before sending data".into());
    }

    let request_str = String::from_utf8_lossy(&peek_buf);

    return Ok(bytes_read == peek_buf.len()
        && request_str.starts_with("GET")
        && !request_str.ends_with("api"));
}

async fn serve_frontend_file(
    connection: NetworkConnection,
) -> Result<(), Box<dyn error::Error + Send + Sync>> {
    let stream = connection.0;

    let app = Router::new()
        .fallback_service(
            ServeDir::new(STATIC_PATH.get().unwrap()).append_index_html_on_directories(true),
        )
        .layer(middleware::from_fn(add_connection_close));

    let tower_svc = app.into_service();
    let hyper_svc = TowerToHyperService::new(tower_svc);

    auto::Builder::new(TokioExecutor::new())
        .serve_connection(TokioIo::new(stream), hyper_svc)
        .await?;

    Ok(())
}

async fn add_connection_close(req: Request<Body>, next: Next) -> Response {
    let mut res = next.run(req).await;

    res.headers_mut()
        .insert(header::CONNECTION, HeaderValue::from_static("close"));
    res
}