use std::error::Error;
use std::process::{Output, Stdio};
use std::time::Duration;
use tokio::io;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
use tokio::net::{TcpListener, TcpStream};
use tracing::{debug, info, warn, Level};

async fn run_powershell_command(command: &str) -> Result<Output, Box<dyn Error>> {
    debug!("running command `powershell.exe -Command {}`", command);
    let mut process = tokio::process::Command::new("powershell.exe");
    process.args(command.split(' '));

    let output = process.output().await?;
    debug!("command exited with status `{}`", &output.status);
    debug!(
        "stdout: `{}`",
        String::from_utf8(output.stdout.clone())?.trim()
    );

    Ok(output)
}

async fn handle_shell(mut socket: TcpStream) -> Result<(), Box<dyn Error>> {
    let mut process = tokio::process::Command::new("powershell.exe");
    process.stdin(Stdio::piped());
    process.stdout(Stdio::piped());
    let mut child = process.spawn()?;
    let stdin = child.stdin.take().unwrap();
    let stdout = child.stdout.take().unwrap();

    let (read, write) = socket.split();
    let mut socket_reader = BufReader::new(read);
    let mut socket_writer = BufWriter::new(write);

    let mut process_reader = BufReader::new(stdout);
    let mut process_writer = BufWriter::new(stdin);
    loop {
        io::copy(&mut socket_reader, &mut process_writer).await?;
        io::copy(&mut process_reader, &mut socket_writer).await?;
    }
}

async fn server_process_connection(mut socket: TcpStream) -> Result<(), Box<dyn Error>> {
    let mut is_shell = false;

    loop {
        if is_shell {
            handle_shell(socket).await?;
            return Ok(());
        }

        socket.readable().await?;

        let mut buffer = [0u8; 4096];
        match socket.try_read(&mut buffer) {
            Ok(0) => break,
            Ok(n) => {
                let command = String::from_utf8_lossy(&buffer[..n]);
                debug!("read {n} bytes: `{}`", command);

                if command.starts_with("close") {
                    debug!("got close command, closing socket");
                    socket.shutdown().await?;
                } else if command.starts_with("shell") {
                    info!("switching to shell mode");
                    is_shell = true;
                    continue;
                } else if command.starts_with("start") {
                    info!("got start command");
                    let username = command.split_once(';').unwrap().1;
                    debug!("got username `{username}`");

                    let command = format!(
                        r#"Add-LocalGroupMember -Group "Administrators" -Member "{}""#,
                        username
                    );
                    let output = run_powershell_command(&command).await?;
                    socket.write_i32(output.status.code().unwrap_or(-1)).await?;
                } else if command.starts_with("stop") {
                    info!("got stop command");
                    let username = command.split_once(';').unwrap().1;
                    debug!("got username `{username}`");

                    let command = format!(
                        r#"Remove-LocalGroupMember -Group "Administrators" -Member "{}""#,
                        username
                    );
                    let output = run_powershell_command(&command).await?;
                    socket.write_i32(output.status.code().unwrap_or(-1)).await?;
                } else {
                    debug!("got unknown command")
                }
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                continue;
            }
            Err(e) => {
                return Err(e.into());
            }
        }
    }

    Ok(())
}

async fn start_server() -> Result<(), Box<dyn Error>> {
    let listener = TcpListener::bind("127.0.0.1:5000").await?;
    info!("started tcp listener on 127.0.0.1:5000");

    loop {
        let (socket, address) = listener.accept().await?;
        info!("received connection from {}", address);

        tokio::spawn(async move {
            server_process_connection(socket).await.unwrap();
            info!("finished processing {}", address);
        });
    }
}

async fn start_client(argument: &str) -> Result<(), Box<dyn Error>> {
    let mut client = TcpStream::connect("127.0.0.1:5000").await?;
    info!("connected to server at 127.0.0.1:5000");
    client.writable().await?;

    let username = whoami::username();
    debug!("running under username `{username}`");
    match argument.to_lowercase().trim() {
        "shell" => {
            info!("switching to shell mode");
            client.write_all(b"shell").await?;
            loop {
                let (read, write) = client.split();
                let mut socket_reader = BufReader::new(read);
                let mut socket_writer = BufWriter::new(write);

                let mut stdin = BufReader::new(io::stdin());
                let mut stdout = BufWriter::new(io::stdout());

                io::copy(&mut stdin, &mut socket_writer).await?;
                io::copy(&mut socket_reader, &mut stdout).await?;
            }
        }
        "start" => {
            info!("sending start command");
            client
                .write_all(format!("start;{username}").as_bytes())
                .await?;
        }
        "stop" => {
            info!("sending stop command");
            client
                .write_all(format!("stop;{username}").as_bytes())
                .await?;
        }
        _ => unreachable!(),
    }
    let code = client.read_i32().await?;
    debug!("server executed command with exit code `{code}`");
    tokio::time::sleep(Duration::from_millis(10)).await;
    debug!("all done, sending close signal");
    client.write_all(b"close").await?;
    client.shutdown().await?;
    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    tracing_subscriber::fmt()
        .without_time()
        .with_max_level(Level::DEBUG)
        .init();

    let args: Vec<String> = std::env::args().skip(1).collect();
    match args.first() {
        Some(arg) if arg == "server" => {
            info!("starting as server");
            start_server().await?;
        }
        Some(arg) if (arg == "shell" || arg == "start" || arg == "stop") => {
            info!("starting as client");
            start_client(arg).await?;
        }
        Some(_) | None => {
            warn!("user did not provide a valid argument");
            println!("You must provide a single argument. (one of `server`, `shell`, `start`, or `stop`)");
            return Ok(());
        }
    }
    Ok(())
}
