mod servercertdb;
mod pendingdataconndb;
mod servermanager;
mod dataconnmanager;

use servercertdb::*;
use pendingdataconndb::*;
use servermanager::*;
use dataconnmanager::*;
use actix::prelude::*;
use std::sync::Arc;
use libbonknet::*;
use rustls::{RootCertStore, ServerConfig};
use rustls::server::WebPkiClientVerifier;
use actix_tls::accept::rustls_0_22::{Acceptor as RustlsAcceptor, TlsStream};
use actix_server::Server;
use actix_rt::net::TcpStream;
use actix_service::{ServiceFactoryExt as _};
use futures::{StreamExt, SinkExt};
use tokio_util::codec::{Framed, FramedRead, FramedWrite, LengthDelimitedCodec};
use tracing::{info, error, warn};
use rcgen::{Certificate, CertificateParams, DnType, KeyPair};
use tokio::io::{ReadHalf, WriteHalf};

type TransportStream = Framed<TlsStream<TcpStream>, LengthDelimitedCodec>;
type TransportStreamTx = FramedWrite<WriteHalf<TlsStream<TcpStream>>, LengthDelimitedCodec>;
type TransportStreamRx = FramedRead<ReadHalf<TlsStream<TcpStream>>, LengthDelimitedCodec>;

struct ServerCert {
    cert: Vec<u8>,
    prkey: Vec<u8>,
}

fn generate_server_cert(root_cert: &Certificate, name: &str) -> ServerCert {
    let mut params = CertificateParams::new(vec!["entity.other.host".into(), format!("bonk.server.{name}")]);
    params.distinguished_name.push(DnType::CommonName, name);
    params.use_authority_key_identifier_extension = true;
    params.key_usages.push(rcgen::KeyUsagePurpose::DigitalSignature);
    params
        .extended_key_usages
        .push(rcgen::ExtendedKeyUsagePurpose::ClientAuth);
    let certificate = Certificate::from_params(params).unwrap();
    let cert = certificate.serialize_der_with_signer(root_cert).unwrap();
    let prkey = certificate.serialize_private_key_der();
    ServerCert { cert, prkey }
}


#[actix_rt::main]
async fn main() {
    // Tracing Subscriber
    let subscriber = tracing_subscriber::FmtSubscriber::new();
    tracing::subscriber::set_global_default(subscriber).unwrap();
    // BROKER CERTS
    let broker_root_cert_der = load_cert("certs/broker_root_cert.pem").unwrap();
    let broker_cert_der = load_cert("certs/broker_cert.pem").unwrap();
    let broker_prkey_der = load_prkey("certs/broker_key.pem").unwrap();
    // SERVER ROOT
    let server_root_cert_der = load_cert("certs/server_root_cert.pem").unwrap();
    let server_root_prkey_der = load_prkey("certs/server_root_key.pem").unwrap();
    // GUESTSERVER ROOT
    let guestserver_root_cert_der = load_cert("certs/guestserver_root_cert.pem").unwrap();
    // CLIENT ROOT
    let client_root_cert_der = load_cert("certs/client_root_cert.pem").unwrap();
    // Client Verifier
    let mut broker_root_store = RootCertStore::empty();
    broker_root_store.add(server_root_cert_der.clone()).unwrap();
    broker_root_store.add(client_root_cert_der.clone()).unwrap();
    broker_root_store.add(guestserver_root_cert_der.clone()).unwrap();
    let server_verifier = WebPkiClientVerifier::builder(Arc::new(broker_root_store)).build().unwrap();
    // Configure TLS
    let server_tlsconfig = ServerConfig::builder()
        // .with_no_client_auth()
        .with_client_cert_verifier(server_verifier)
        .with_single_cert(vec![broker_cert_der.clone(), broker_root_cert_der.clone()], broker_prkey_der.into())
        .unwrap();
    let server_acceptor = RustlsAcceptor::new(server_tlsconfig);

    let server_root_cert_der = Arc::new(server_root_cert_der);
    let server_root_prkey = KeyPair::from_der(server_root_prkey_der.secret_pkcs8_der()).unwrap();
    let client_root_cert_der = Arc::new(client_root_cert_der);
    let guestserver_root_cert_der = Arc::new(guestserver_root_cert_der);
    let server_root_cert = Arc::new(Certificate::from_params(CertificateParams::from_ca_cert_der(
        &server_root_cert_der,
        server_root_prkey
    ).unwrap()).unwrap());

    let scdb_addr = ServerCertDB::new().start();
    let dcm_addr = DataConnManager::new().start();
    let pdcm_addr = PendingDataConnManager::new(dcm_addr.clone()).start();
    let sm_addr = ServerManager::new(pdcm_addr.clone()).start();

    Server::build()
        .bind("server-command", ("localhost", 2541), move || {
            let server_root_cert_der = Arc::clone(&server_root_cert_der);
            let client_root_cert_der = Arc::clone(&client_root_cert_der);
            let guestserver_root_cert_der = Arc::clone(&guestserver_root_cert_der);
            let server_root_cert = Arc::clone(&server_root_cert);
            let scdb_addr = scdb_addr.clone();
            let pdcm_addr = pdcm_addr.clone();
            let sm_addr = sm_addr.clone();

            // Set up TLS service factory
            server_acceptor
                .clone()
                .map_err(|err| println!("Rustls error: {:?}", err))
                .and_then(move |stream: TlsStream<TcpStream>| {
                    let server_root_cert_der = Arc::clone(&server_root_cert_der);
                    let client_root_cert_der = Arc::clone(&client_root_cert_der);
                    let guestserver_root_cert_der = Arc::clone(&guestserver_root_cert_der);
                    let server_root_cert = Arc::clone(&server_root_cert);
                    let scdb_addr = scdb_addr.clone();
                    let pdcm_addr = pdcm_addr.clone();
                    let sm_addr = sm_addr.clone();
                    async move {
                        let peer_certs = stream.get_ref().1.peer_certificates().unwrap();
                        let peer_cert_bytes = peer_certs.first().unwrap().to_vec();
                        let peer_root_cert_der = peer_certs.last().unwrap().clone();
                        if peer_root_cert_der == *server_root_cert_der {
                            info!("Server connection");
                            let mut transport = Framed::new(stream, LengthDelimitedCodec::new());
                            match transport.next().await {
                                None => {
                                    info!("Connection closed by peer");
                                }
                                Some(item) => match item {
                                    Ok(buf) => {
                                        use FromServerConnTypeMessage::*;
                                        let msg: FromServerConnTypeMessage = rmp_serde::from_slice(&buf).unwrap();
                                        info!("{:?}", msg);
                                        match msg {
                                            SendCommand => {
                                                info!("SendCommand Stream");
                                                let reply = ToServerConnTypeReply::OkSendCommand;
                                                transport.send(rmp_serde::to_vec(&reply).unwrap().into()).await.unwrap();
                                                server_command_handler(transport, peer_cert_bytes, scdb_addr).await;
                                            }
                                            Subscribe => {
                                                info!("Subscribe Stream");
                                                let name = match scdb_addr.send(FetchName { cert: peer_cert_bytes }).await.unwrap() {
                                                    None => {
                                                        error!("Cert has no name assigned!");
                                                        let reply = ToServerConnTypeReply::GenericFailure;
                                                        transport.send(rmp_serde::to_vec(&reply).unwrap().into()).await.unwrap();
                                                        return Ok(());
                                                    }
                                                    Some(name) => name,
                                                };
                                                let reply = ToServerConnTypeReply::OkSubscribe;
                                                transport.send(rmp_serde::to_vec(&reply).unwrap().into()).await.unwrap();
                                                server_subscribe_handler(transport, name, sm_addr).await;
                                            }
                                            OpenDataStream(conn_id) => {
                                                info!("OpenDataStream with {:?}", conn_id);
                                                let msg = RegisterStream::server(conn_id, transport);
                                                pdcm_addr.send(msg).await.unwrap().unwrap();
                                            }
                                        }
                                    }
                                    Err(e) => {
                                        info!("Disconnection: {:?}", e);
                                    }
                                }
                            }
                            info!("Server Task terminated!");
                        } else if peer_root_cert_der == *guestserver_root_cert_der {
                            info!("GuestServer connection");
                            let server_root_cert = Arc::clone(&server_root_cert);
                            let codec = LengthDelimitedCodec::new();
                            let transport = Framed::new(stream, codec);
                            guestserver_handler(transport, scdb_addr, &server_root_cert).await;
                        } else if peer_root_cert_der == *client_root_cert_der {
                            info!("Client connection");
                            let codec = LengthDelimitedCodec::new();
                            let transport = Framed::new(stream, codec);
                            client_handler(transport, sm_addr, pdcm_addr).await;
                        } else {
                            error!("Unknown Root Certificate");
                        }
                        Ok(())
                    }
                })
        }).unwrap()
        .workers(1)
        .run()
        .await
        .unwrap();
}

async fn server_subscribe_handler(transport: TransportStream, name: String, sm_addr: Addr<ServerManager>) {
    match sm_addr.send(StartTransporter { name, transport }).await.unwrap() {
        Ok(_) => {
            info!("Stream sent to the manager");
        }
        Err(e) => {
            error!("Error from manager: {:?}", e);
        }
    }
}

async fn server_command_handler(mut transport: TransportStream, peer_cert_bytes: Vec<u8>, server_db_addr: Addr<ServerCertDB>) {
    loop {
        match transport.next().await {
            None => {
                info!("Transport returned None");
                break;
            }
            Some(item) => match item {
                Ok(buf) => {
                    use FromServerCommandMessage::*;
                    let msg: FromServerCommandMessage = rmp_serde::from_slice(&buf).unwrap();
                    info!("{:?}", msg);
                    match msg {
                        ChangeName { name } => {
                            info!("Changing name to {}", name);
                            match server_db_addr.send(UnregisterServer { cert: peer_cert_bytes.clone() }).await.unwrap() {
                                None => {
                                    info!("Nothing to unregister");
                                }
                                Some(old_name) => {
                                    warn!("Unregistered from old name {}", old_name);
                                }
                            }
                            let reply = match server_db_addr.send(RegisterServer { cert: peer_cert_bytes.clone(), name }).await.unwrap() {
                                Ok(_) => {
                                    info!("Registered!");
                                    ToServerCommandReply::NameChanged
                                }
                                Err(e) => {
                                    error!("Error registering: {:?}", e);
                                    ToServerCommandReply::GenericFailure
                                }
                            };
                            transport.send(rmp_serde::to_vec(&reply).unwrap().into()).await.unwrap();
                        }
                        WhoAmI => {
                            info!("Asked who I am");
                            let reply = match server_db_addr.send(FetchName { cert: peer_cert_bytes.clone() }).await.unwrap() {
                                None => {
                                    info!("I'm not registered anymore!? WTF");
                                    ToServerCommandReply::GenericFailure
                                }
                                Some(name) => {
                                    info!("I am {}", name);
                                    ToServerCommandReply::YouAre { name }
                                }
                            };
                            transport.send(rmp_serde::to_vec(&reply).unwrap().into()).await.unwrap();
                        }
                    }
                }
                Err(e) => {
                    info!("Disconnection: {:?}", e);
                    break;
                }
            }
        }
    }
}

// TODO: Considera creare un context dove vengono contenute tutte le chiavi e gli address da dare a tutti gli handler
async fn guestserver_handler(mut transport: TransportStream, server_db_addr: Addr<ServerCertDB>, server_root_cert: &Arc<Certificate>) {
    loop {
        match transport.next().await {
            None => {
                info!("Transport returned None");
                break;
            }
            Some(item) => {
                match item {
                    Ok(buf) => {
                        use FromGuestServerMessage::*;
                        let msg: FromGuestServerMessage = rmp_serde::from_slice(&buf).unwrap();
                        info!("{:?}", msg);
                        match msg {
                            Announce { name } => {
                                info!("Announced with name {}", name);
                                if server_db_addr.send(IsNameRegistered { name: name.clone() }).await.unwrap() {
                                    info!("Name {} already registered!", name);
                                    let reply = ToGuestServerMessage::FailedNameAlreadyOccupied;
                                    transport.send(rmp_serde::to_vec(&reply).unwrap().into()).await.unwrap();
                                    break; // Stop reading
                                } else {
                                    let cert = generate_server_cert(server_root_cert, name.as_str());
                                    server_db_addr.send(RegisterServer {
                                        cert: cert.cert.clone(),
                                        name,
                                    }).await.unwrap().unwrap();
                                    let reply = ToGuestServerMessage::OkAnnounce {
                                        server_cert: cert.cert,
                                        server_prkey: cert.prkey
                                    };
                                    transport.send(rmp_serde::to_vec(&reply).unwrap().into()).await.unwrap();
                                }
                            }
                        }
                    }
                    Err(e) => {
                        info!("Disconnection: {:?}", e);
                        break;
                    }
                }
            }
        }
    }
}

async fn client_handler(mut transport: TransportStream, sm_addr: Addr<ServerManager>, pdcm_addr: Addr<PendingDataConnManager>) {
    loop {
        match transport.next().await {
            None => {
                info!("Transport returned None");
                break;
            }
            Some(item) => {
                match item {
                    Ok(buf) => {
                        let msg: FromClientCommand = rmp_serde::from_slice(&buf).unwrap();
                        info!("{:?}", msg);
                        match msg {
                            FromClientCommand::RequestServer { name } => {
                                info!("REQUESTED SERVER {}", name);
                                let data = sm_addr.send(RequestServer { name }).await.unwrap();
                                match data {
                                    Ok(client_conn_id) => {
                                        let reply = ToClientResponse::OkRequest { conn_id: client_conn_id };
                                        transport.send(rmp_serde::to_vec(&reply).unwrap().into()).await.unwrap();
                                        info!("Sent OkRequest");
                                    }
                                    Err(e) => {
                                        error!("Error! {:?}", e);
                                    }
                                }
                            }
                            FromClientCommand::ServerList => {
                                info!("Requested ServerList");
                                let data = sm_addr.send(GetServerList {}).await.unwrap();
                                let reply = ToClientResponse::OkServerList { data };
                                transport.send(rmp_serde::to_vec(&reply).unwrap().into()).await.unwrap();
                            }
                            FromClientCommand::UpgradeToDataStream(conn_id) => {
                                info!("Upgrade to DataStream with conn_id {:?}", conn_id);
                                let msg = RegisterStream::client(conn_id, transport);
                                pdcm_addr.send(msg).await.unwrap().unwrap();
                                break;
                            }
                        }
                    }
                    Err(e) => {
                        info!("Disconnection: {:?}", e);
                        break;
                    }
                }
            }
        }
    }
}