//! Windows SYSTEM service for UNATTENDED access + lock/login-screen capture.
//!
//! The service runs as LocalSystem (session 0) and spawns the NetDesk host INTO
//! the active console session **as SYSTEM**, with its thread desktop following
//! the current input desktop. That lets the host capture and inject on the
//! secure desktop (the lock screen, the Windows sign-in screen, and UAC) — so a
//! remote viewer can sign in and unlock an unattended machine, TeamViewer-style.
//!
//! Opt-in and isolated: nothing here runs unless `--install-service` is executed
//! elevated. The ordinary user-session host (Desktop autostart / admin task) is
//! untouched, so installing or removing the service never risks the path that
//! already works.
//!
//! Entry points (wired in the viewer binary's `main`):
//!   --install-service    register + start the service (self-elevates via UAC)
//!   --uninstall-service  stop + delete the service
//!   --service-run        the SCM service entry (run by Windows, not by hand)
//!   --session-agent      the per-session host the service spawns (follows the
//!                        input desktop so the lock screen is captured)

#![cfg(windows)]

use anyhow::{anyhow, Result};
use std::ffi::c_void;
use std::sync::atomic::{AtomicBool, Ordering};
use windows::core::{PCWSTR, PWSTR};
use windows::Win32::Foundation::{CloseHandle, HANDLE, LUID};
use windows::Win32::Security::{
    AdjustTokenPrivileges, DuplicateTokenEx, LookupPrivilegeValueW, SetTokenInformation,
    SecurityIdentification, TokenPrimary, TokenSessionId, LUID_AND_ATTRIBUTES,
    SE_PRIVILEGE_ENABLED, TOKEN_ACCESS_MASK, TOKEN_ADJUST_DEFAULT, TOKEN_ADJUST_PRIVILEGES,
    TOKEN_ADJUST_SESSIONID, TOKEN_ASSIGN_PRIMARY, TOKEN_DUPLICATE, TOKEN_PRIVILEGES, TOKEN_QUERY,
};
use windows::Win32::System::RemoteDesktop::WTSGetActiveConsoleSessionId;
use windows::Win32::System::Services::{
    RegisterServiceCtrlHandlerW, SetServiceStatus, StartServiceCtrlDispatcherW,
    SERVICE_STATUS, SERVICE_STATUS_HANDLE, SERVICE_TABLE_ENTRYW, SERVICE_ACCEPT_STOP,
    SERVICE_CONTROL_SHUTDOWN, SERVICE_CONTROL_STOP, SERVICE_RUNNING, SERVICE_START_PENDING,
    SERVICE_STOP_PENDING, SERVICE_STOPPED, SERVICE_WIN32_OWN_PROCESS,
};
use windows::Win32::System::Threading::{
    CreateProcessAsUserW, GetCurrentProcess, OpenProcessToken, TerminateProcess,
    CREATE_NO_WINDOW, CREATE_UNICODE_ENVIRONMENT, PROCESS_INFORMATION, STARTUPINFOW,
};

const SERVICE_NAME: &str = "NetDeskHost";

static STOP: AtomicBool = AtomicBool::new(false);
static mut STATUS_HANDLE: isize = 0;

fn wide(s: &str) -> Vec<u16> {
    s.encode_utf16().chain(std::iter::once(0)).collect()
}

// ---------------------------------------------------------------------------
// install / uninstall (shell out to sc.exe; requires elevation)
// ---------------------------------------------------------------------------

/// Register the service (LocalSystem, auto-start, auto-restart) and start it.
pub fn install_service(exe: &str) -> Result<()> {
    let bin = format!("\"{exe}\" --service-run");
    // Create (quirk: sc wants "key=" and the value as separate tokens).
    run_sc(&["create", SERVICE_NAME, "binPath=", &bin, "start=", "auto",
             "obj=", "LocalSystem", "DisplayName=", "NetDesk Host (unattended)"])?;
    let _ = run_sc(&["description", SERVICE_NAME,
        "NetDesk unattended host — captures the lock/sign-in screen so you can access this PC remotely."]);
    // Restart on crash so unattended access survives faults.
    let _ = run_sc(&["failure", SERVICE_NAME, "reset=", "0", "actions=",
        "restart/5000/restart/5000/restart/5000"]);
    run_sc(&["start", SERVICE_NAME])?;
    Ok(())
}

/// Stop and delete the service.
pub fn uninstall_service() -> Result<()> {
    let _ = run_sc(&["stop", SERVICE_NAME]);
    run_sc(&["delete", SERVICE_NAME])?;
    Ok(())
}

fn run_sc(args: &[&str]) -> Result<()> {
    let status = std::process::Command::new("sc")
        .args(args)
        .status()
        .map_err(|e| anyhow!("sc {args:?}: {e}"))?;
    if status.success() {
        Ok(())
    } else {
        Err(anyhow!("sc {:?} exited with {:?}", args, status.code()))
    }
}

// ---------------------------------------------------------------------------
// service entry (run by the SCM via `--service-run`)
// ---------------------------------------------------------------------------

/// Hand control to the Service Control Manager. Blocks until the service stops.
pub fn run_service() -> Result<()> {
    let name = wide(SERVICE_NAME);
    let table = [
        SERVICE_TABLE_ENTRYW {
            lpServiceName: PWSTR(name.as_ptr() as *mut u16),
            lpServiceProc: Some(service_main),
        },
        SERVICE_TABLE_ENTRYW {
            lpServiceName: PWSTR::null(),
            lpServiceProc: None,
        },
    ];
    unsafe {
        StartServiceCtrlDispatcherW(table.as_ptr())
            .map_err(|e| anyhow!("StartServiceCtrlDispatcher: {e}"))?;
    }
    Ok(())
}

extern "system" fn service_main(_argc: u32, _argv: *mut PWSTR) {
    unsafe {
        let name = wide(SERVICE_NAME);
        match RegisterServiceCtrlHandlerW(PCWSTR(name.as_ptr()), Some(handler)) {
            Ok(h) => STATUS_HANDLE = h.0 as isize,
            Err(_) => return,
        }
        set_status(SERVICE_START_PENDING, 0);
        set_status(SERVICE_RUNNING, SERVICE_ACCEPT_STOP);
        // Main loop: keep a host alive in the active console session.
        let exe = std::env::current_exe()
            .map(|p| p.display().to_string())
            .unwrap_or_default();
        let mut child: Option<(HANDLE, u32)> = None;
        while !STOP.load(Ordering::Relaxed) {
            let session = WTSGetActiveConsoleSessionId();
            let need_spawn = match child {
                // 0xFFFFFFFF => no one attached to the console right now.
                _ if session == 0xFFFF_FFFF => false,
                None => true,
                Some((_, sess)) => sess != session || !process_alive(child),
            };
            if need_spawn && session != 0xFFFF_FFFF {
                if let Some((h, _)) = child.take() {
                    let _ = TerminateProcess(h, 1);
                    let _ = CloseHandle(h);
                }
                match spawn_host_in_session(session, &exe) {
                    Ok(h) => child = Some((h, session)),
                    Err(_) => {}
                }
            }
            std::thread::sleep(std::time::Duration::from_millis(1500));
        }
        if let Some((h, _)) = child.take() {
            let _ = TerminateProcess(h, 0);
            let _ = CloseHandle(h);
        }
        set_status(SERVICE_STOPPED, 0);
    }
}

extern "system" fn handler(control: u32) {
    if control == SERVICE_CONTROL_STOP || control == SERVICE_CONTROL_SHUTDOWN {
        STOP.store(true, Ordering::Relaxed);
        unsafe { set_status(SERVICE_STOP_PENDING, 0) };
    }
}

unsafe fn set_status(state: windows::Win32::System::Services::SERVICE_STATUS_CURRENT_STATE, accept: u32) {
    if STATUS_HANDLE == 0 {
        return;
    }
    let mut st = SERVICE_STATUS {
        dwServiceType: SERVICE_WIN32_OWN_PROCESS,
        dwCurrentState: state,
        dwControlsAccepted: accept,
        dwWin32ExitCode: 0,
        dwServiceSpecificExitCode: 0,
        dwCheckPoint: 0,
        dwWaitHint: 3000,
    };
    let _ = SetServiceStatus(SERVICE_STATUS_HANDLE(STATUS_HANDLE as *mut c_void), &mut st);
}

fn process_alive(child: Option<(HANDLE, u32)>) -> bool {
    use windows::Win32::Foundation::{STILL_ACTIVE, WAIT_OBJECT_0};
    use windows::Win32::System::Threading::{GetExitCodeProcess, WaitForSingleObject};
    match child {
        None => false,
        Some((h, _)) => unsafe {
            // Fast path: a 0ms wait signals only if the process has exited.
            if WaitForSingleObject(h, 0) == WAIT_OBJECT_0 {
                return false;
            }
            let mut code = 0u32;
            if GetExitCodeProcess(h, &mut code).is_ok() {
                code == STILL_ACTIVE.0 as u32
            } else {
                false
            }
        },
    }
}

// ---------------------------------------------------------------------------
// spawn the host as SYSTEM in the active session (so it can reach the secure
// desktop). Duplicates the service's SYSTEM token, re-targets it at the active
// session, and CreateProcessAsUser's the host on winsta0\default.
// ---------------------------------------------------------------------------

fn spawn_host_in_session(session: u32, exe: &str) -> Result<HANDLE> {
    unsafe {
        let proc = GetCurrentProcess();
        let access = TOKEN_ACCESS_MASK(
            TOKEN_DUPLICATE.0
                | TOKEN_QUERY.0
                | TOKEN_ASSIGN_PRIMARY.0
                | TOKEN_ADJUST_DEFAULT.0
                | TOKEN_ADJUST_SESSIONID.0
                | TOKEN_ADJUST_PRIVILEGES.0,
        );
        let mut tok = HANDLE::default();
        OpenProcessToken(proc, access, &mut tok).map_err(|e| anyhow!("OpenProcessToken: {e}"))?;
        enable_privilege(tok, "SeTcbPrivilege");

        let mut dup = HANDLE::default();
        DuplicateTokenEx(
            tok,
            TOKEN_ACCESS_MASK(0x02000000), // MAXIMUM_ALLOWED
            None,
            SecurityIdentification,
            TokenPrimary,
            &mut dup,
        )
        .map_err(|e| anyhow!("DuplicateTokenEx: {e}"))?;

        // Re-target the primary token at the active console session.
        let sess = session;
        SetTokenInformation(
            dup,
            TokenSessionId,
            &sess as *const u32 as *const c_void,
            std::mem::size_of::<u32>() as u32,
        )
        .map_err(|e| anyhow!("SetTokenInformation(session): {e}"))?;

        let mut cmd = wide(&format!("\"{exe}\" --session-agent"));
        let mut desktop = wide("winsta0\\default");
        let mut si = STARTUPINFOW::default();
        si.cb = std::mem::size_of::<STARTUPINFOW>() as u32;
        si.lpDesktop = PWSTR(desktop.as_mut_ptr());
        let mut pi = PROCESS_INFORMATION::default();

        CreateProcessAsUserW(
            dup,
            PCWSTR::null(),
            PWSTR(cmd.as_mut_ptr()),
            None,
            None,
            false,
            CREATE_NO_WINDOW | CREATE_UNICODE_ENVIRONMENT,
            None,
            PCWSTR::null(),
            &si,
            &mut pi,
        )
        .map_err(|e| anyhow!("CreateProcessAsUser: {e}"))?;

        let _ = CloseHandle(pi.hThread);
        let _ = CloseHandle(dup);
        let _ = CloseHandle(tok);
        Ok(pi.hProcess)
    }
}

fn enable_privilege(token: HANDLE, name: &str) {
    unsafe {
        let wname = wide(name);
        let mut luid = LUID::default();
        if LookupPrivilegeValueW(PCWSTR::null(), PCWSTR(wname.as_ptr()), &mut luid).is_err() {
            return;
        }
        let tp = TOKEN_PRIVILEGES {
            PrivilegeCount: 1,
            Privileges: [LUID_AND_ATTRIBUTES {
                Luid: luid,
                Attributes: SE_PRIVILEGE_ENABLED,
            }],
        };
        let _ = AdjustTokenPrivileges(token, false, Some(&tp), 0, None, None);
    }
}