use std::error::Error;
use std::process::Output;
use std::time::Duration;
use tokio::io;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tracing::level_filters::LevelFilter;
#[allow(unused_imports)]
use tracing::{debug, info, Level};
use tracing_subscriber::EnvFilter;

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 server_process_connection(mut socket: TcpStream) -> Result<(), Box<dyn Error>> {
    loop {
        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");
                    let command = r#"Start-Process powershell.exe"#;
                    let output = run_powershell_command(command).await?;
                    socket.write_i32(output.status.code().unwrap_or(-1)).await?;
                } 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}`");

    let mut restart_explorer = false;
    match argument.to_lowercase().trim() {
        "shell" => {
            info!("opening shell");
            client.write_all(b"shell").await?;
        }
        "start" => {
            info!("sending start command");
            client
                .write_all(format!("start;{username}").as_bytes())
                .await?;
            restart_explorer = true;
        }
        "stop" => {
            info!("sending stop command");
            client
                .write_all(format!("stop;{username}").as_bytes())
                .await?;
            restart_explorer = true;
        }
        _ => unreachable!(),
    }
    let code = client.read_i32().await?;
    debug!("server executed command with exit code `{code}`");

    tokio::time::sleep(Duration::from_millis(100)).await;
    if code == 0 && restart_explorer {
        info!("restarting explorer");
        run_powershell_command(r#"Stop-Process -ProcessName explorer -Force"#).await?;
    }
    debug!("all done, sending close signal");
    client.write_all(b"close").await?;
    client.shutdown().await?;

    info!("done");
    Ok(())
}

fn init_logging(as_server: bool) {
    let filter = EnvFilter::builder()
        .with_default_directive(LevelFilter::DEBUG.into())
        .from_env_lossy();
    #[allow(unused_variables)]
    let appender = tracing_appender::rolling::never(std::env::temp_dir(), "superman.log");
    if as_server {
        #[cfg(not(debug_assertions))]
        tracing_subscriber::fmt()
            .with_max_level(Level::TRACE)
            .with_ansi(false)
            .with_writer(appender)
            .init();
        #[cfg(debug_assertions)]
        tracing_subscriber::fmt()
            .with_env_filter(filter)
            .with_writer(std::io::stdout)
            .init();
    } else {
        tracing_subscriber::fmt()
            .with_env_filter(filter)
            .without_time()
            .with_writer(std::io::stdout)
            .init();
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let args: Vec<String> = std::env::args().skip(1).collect();
    match args.first() {
        Some(arg) if arg == "server" => {
            init_logging(true);
            info!("starting as server");
            start_server().await?;
        }
        Some(arg) if (arg == "shell" || arg == "start" || arg == "stop") => {
            init_logging(false);
            info!("starting as client");
            start_client(arg).await?;
        }
        Some(_) | None => {
            println!("You must provide a single argument. (one of `server`, `shell`, `start`, or `stop`)");
            return Ok(());
        }
    }
    Ok(())
}
