use rshim::Shim;
use std::os::windows::ffi::OsStrExt;
use std::{env, ffi::CString, mem::size_of, path::Path, process::exit, ptr::null_mut};
use tracing_subscriber::{EnvFilter, fmt};

use std::ffi::{OsStr, OsString};

use windows_sys::Win32::UI::Shell::{SHELLEXECUTEINFOW, ShellExecuteExW};

use windows_sys::Win32::Foundation::CloseHandle;
use windows_sys::{
    Win32::{
        Foundation::{FALSE, TRUE},
        System::{
            Com::{COINIT_APARTMENTTHREADED, COINIT_DISABLE_OLE1DDE, CoInitializeEx},
            Console::{
                CTRL_BREAK_EVENT, CTRL_C_EVENT, CTRL_CLOSE_EVENT, CTRL_LOGOFF_EVENT,
                CTRL_SHUTDOWN_EVENT, SetConsoleCtrlHandler,
            },
            Threading::{GetExitCodeProcess, INFINITE, WaitForSingleObject},
        },
        UI::{
            Shell::{
                SEE_MASK_NOASYNC, SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOA, ShellExecuteExA,
            },
            WindowsAndMessaging::SW_NORMAL,
        },
    },
    core::BOOL,
};

unsafe extern "system" fn console_ctrl_handler(evt: u32) -> BOOL {
    match evt {
        CTRL_C_EVENT => TRUE,        //eprintln!("ctrl_c handled!"),
        CTRL_BREAK_EVENT => TRUE,    //eprintln!("ctrl_break handled!"),
        CTRL_CLOSE_EVENT => TRUE,    //eprintln!("ctrl_close handled!"),
        CTRL_LOGOFF_EVENT => TRUE,   //eprintln!("ctrl_logoff handled!"),
        CTRL_SHUTDOWN_EVENT => TRUE, //eprintln!("ctrl_shutdown handled!"),
        other => {
            eprintln!("未知的系统事件编号: {}，未处理！", other);
            FALSE
        }
    }
}

const EXIT_FAILED_LOAD_SHIM: i32 = 1;
const EXIT_FAILED_SPAWN_PROG: i32 = 2;
const EXIT_FAILED_WAIT_PROG: i32 = 3;
const EXIT_PROG_TERMINATED: i32 = 4;

const ERROR_ELEVATION_REQUIRED: i32 = 740;
fn main() {
    // 初始化日志：输出到 stderr，避免污染 shim 子进程的 stdout
    fmt()
        .with_writer(std::io::stderr)
        .with_env_filter(
            EnvFilter::try_from_env("SHIM_LOG").unwrap_or_else(|_| EnvFilter::new("warn")),
        )
        .init();

    let res: BOOL = unsafe { SetConsoleCtrlHandler(Some(console_ctrl_handler), TRUE) };
    if res == FALSE {
        eprintln!("警告: 注册控制台中断事件处理器失败。");
    }

    let calling_args: Vec<_> = env::args_os().skip(1).collect();
    let shim = match Shim::load() {
        Ok(v) => v,
        Err(e) => {
            eprintln!("加载代理(shim)配置时发生错误: {}", e);
            exit(EXIT_FAILED_LOAD_SHIM);
        }
    };

    // 构建 Command：复用 ShimConfig::to_command（含 args/env 注入），避免重复逻辑
    let mut cmd = shim.to_command(&calling_args);

    // 提权回退时需要完整参数：配置默认参数 + 调用方透传参数
    let mut args = shim.args.clone();
    args.extend_from_slice(&calling_args);

    let mut child = match cmd.spawn() {
        Ok(v) => v,
        Err(e) if e.raw_os_error() == Some(ERROR_ELEVATION_REQUIRED) => {
            exit(execute_elevated(&shim.target, &args, Some(&shim.env)))
        }
        Err(e) => {
            eprintln!(
                "启动目标程序 [{}] 时发生错误: {}",
                shim.target.to_string_lossy(),
                e
            );
            exit(EXIT_FAILED_SPAWN_PROG);
        }
    };
    let status = match child.wait() {
        Ok(v) => v,
        Err(e) => {
            eprintln!(
                "等待目标程序 [{}] 执行完毕时发生错误: {}",
                shim.target.to_string_lossy(),
                e
            );
            exit(EXIT_FAILED_WAIT_PROG);
        }
    };
    exit(status.code().unwrap_or(EXIT_PROG_TERMINATED))
}
// 辅助函数：将任意 OsStr 转换为以 \0 结尾的 UTF-16 宽字符向量 (Vec<u16>)
fn to_wide_null(s: impl AsRef<OsStr>) -> Vec<u16> {
    s.as_ref().encode_wide().chain(std::iter::once(0)).collect()
}

fn execute_elevated(
    program: &Path,
    args: &[OsString],
    env_vars: Option<&std::collections::HashMap<String, OsString>>,
) -> i32 {
    // 若提权启动，在此处将环境变量设置给当前进程（即将弹窗 UAC 的进程，随后会被子进程继承）
    if let Some(env_map) = env_vars {
        for (k, v) in env_map {
            unsafe {
                env::set_var(k, v);
            }
        }
    }
    // 2. 将参数列表按 Windows 命令行规则拼装为单个命令行字符串
    let mut arguments_os = OsString::new();
    for (i, arg) in args.iter().enumerate() {
        if i > 0 {
            arguments_os.push(" ");
        }
        let arg_str = arg.to_string_lossy();
        if arg_str.is_empty() {
            arguments_os.push("\"\"");
        } else if !arg_str.contains([' ', '\t', '"']) {
            arguments_os.push(arg);
        } else {
            // 包含空格或引号时进行标准转义包裹
            arguments_os.push("\"");
            for c in arg_str.chars() {
                match c {
                    '\\' => arguments_os.push("\\\\"),
                    '"' => arguments_os.push("\\\""),
                    _ => arguments_os.push(c.to_string()),
                }
            }
            arguments_os.push("\"");
        }
    }
    // let runas = CString::new("runas").unwrap();
    // let program = CString::new(program.to_str().unwrap()).unwrap();
    // let mut arguments = String::new();
    // for arg in args.iter() {
    //     arguments.push(' ');
    //     if arg.len() == 0 {
    //         arguments.push_str("\"\"");
    //     } else if arg.find(&[' ', '\t', '"'][..]).is_none() {
    //         arguments.push_str(&arg);
    //     } else {
    //         arguments.push('"');
    //         for c in arg.chars() {
    //             match c {
    //                 '\\' => arguments.push_str("\\\\"),
    //                 '"' => arguments.push_str("\\\""),
    //                 c => arguments.push(c),
    //             }
    //         }
    //         arguments.push('"');
    //     }
    // }
    // 3. 准备 Windows 宽字符参数
    let runas = to_wide_null("runas");
    let program_wide = to_wide_null(program.as_os_str());
    let arguments_wide = to_wide_null(&arguments_os);

    let mut info = SHELLEXECUTEINFOW {
        cbSize: size_of::<SHELLEXECUTEINFOW>() as u32,
        fMask: SEE_MASK_NOASYNC | SEE_MASK_NOCLOSEPROCESS,
        hwnd: null_mut(),
        lpVerb: runas.as_ptr(),
        lpFile: program_wide.as_ptr(),
        lpParameters: arguments_wide.as_ptr(),
        lpDirectory: null_mut(),
        nShow: SW_NORMAL as i32,
        hInstApp: null_mut(),
        lpIDList: null_mut(),
        lpClass: null_mut(),
        hkeyClass: null_mut(),
        dwHotKey: 0,
        Anonymous: unsafe { std::mem::zeroed() },
        hProcess: null_mut(),
    };
    // let arguments = CString::new(&arguments[..]).unwrap();
    // let mut info = SHELLEXECUTEINFOA::default();
    // info.cbSize = size_of::<SHELLEXECUTEINFOA>() as u32;
    // info.fMask = SEE_MASK_NOASYNC | SEE_MASK_NOCLOSEPROCESS;
    // info.lpVerb = runas.as_ptr().cast::<u8>();
    // info.lpFile = program.as_ptr().cast::<u8>();
    // info.lpParameters = arguments.as_ptr().cast::<u8>();
    // info.nShow = SW_NORMAL;
    let res = unsafe {
        CoInitializeEx(
            null_mut(),
            (COINIT_APARTMENTTHREADED | COINIT_DISABLE_OLE1DDE) as u32,
        );
        // ShellExecuteExA(&mut info as *mut _)
        ShellExecuteExW(&mut info)
    };
    if res == FALSE || info.hProcess == null_mut() {
        return EXIT_FAILED_SPAWN_PROG;
    }
    // 5. 等待提权子进程执行结束并获取退出状态码
    let mut exit_code: u32 = 0;
    unsafe {
        WaitForSingleObject(info.hProcess, INFINITE);
        let ok = GetExitCodeProcess(info.hProcess, &mut exit_code);
        CloseHandle(info.hProcess); // 释放进程句柄，防止资源泄露

        if ok == FALSE {
            return EXIT_FAILED_WAIT_PROG;
        }
    }

    exit_code as i32

    // let mut code: u32 = 0;
    // unsafe {
    //     WaitForSingleObject(info.hProcess, INFINITE);
    //     if GetExitCodeProcess(info.hProcess, &mut code as *mut _) == FALSE {
    //         return EXIT_FAILED_WAIT_PROG;
    //     }
    // }
    // code as i32
}