//! USB Storage Control Module
//!
//! Este modulo implementa o controle de dispositivos de armazenamento USB no Windows.
//! Utiliza duas abordagens complementares:
//! 1. Removable Storage Access Policy (via registro do Windows)
//! 2. USBSTOR driver control (como fallback/reforco)
//!
//! IMPORTANTE: Requer privilegios de administrador para funcionar.

use serde::{Deserialize, Serialize};
use std::io;
use thiserror::Error;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum UsbPolicy {
    Allow,
    BlockAll,
    Readonly,
}

impl UsbPolicy {
    pub fn from_str(s: &str) -> Option<Self> {
        match s.to_uppercase().as_str() {
            "ALLOW" => Some(Self::Allow),
            "BLOCK_ALL" => Some(Self::BlockAll),
            "READONLY" => Some(Self::Readonly),
            _ => None,
        }
    }

    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Allow => "ALLOW",
            Self::BlockAll => "BLOCK_ALL",
            Self::Readonly => "READONLY",
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UsbPolicyResult {
    pub success: bool,
    pub policy: String,
    pub error: Option<String>,
    pub applied_at: Option<i64>,
}

#[derive(Error, Debug)]
#[allow(dead_code)]
pub enum UsbControlError {
    #[error("Politica USB invalida: {0}")]
    InvalidPolicy(String),
    #[error("Erro de registro do Windows: {0}")]
    RegistryError(String),
    #[error("Permissao negada - requer privilegios de administrador")]
    PermissionDenied,
    #[error("Sistema operacional nao suportado")]
    UnsupportedOs,
    #[error("Erro de I/O: {0}")]
    Io(#[from] io::Error),
}

#[cfg(target_os = "windows")]
mod windows_impl {
    use super::*;
    use std::fs;
    use std::path::PathBuf;
    use std::process::Command;
    use winreg::enums::*;
    use winreg::RegKey;

    // GUID para Removable Storage Devices (Disk)
    const REMOVABLE_STORAGE_GUID: &str = "{53f56307-b6bf-11d0-94f2-00a0c91efb8b}";

    // Chaves de registro
    const REMOVABLE_STORAGE_PATH: &str =
        r"Software\Policies\Microsoft\Windows\RemovableStorageDevices";
    const USBSTOR_PATH: &str = r"SYSTEM\CurrentControlSet\Services\USBSTOR";
    const STORAGE_POLICY_PATH: &str = r"SYSTEM\CurrentControlSet\Control\StorageDevicePolicies";

    pub fn apply_usb_policy(policy: UsbPolicy) -> Result<UsbPolicyResult, UsbControlError> {
        let now = chrono::Utc::now().timestamp_millis();

        let direct_result = try_apply_policy_direct(policy);

        match direct_result {
            Ok(()) => Ok(UsbPolicyResult {
                success: true,
                policy: policy.as_str().to_string(),
                error: None,
                applied_at: Some(now),
            }),
            Err(err) => {
                // Se faltou permissão, retorna erro - o serviço deve ser usado
                // Não fazemos elevação aqui para evitar UAC adicional
                if is_permission_error(&err) {
                    return Err(UsbControlError::PermissionDenied);
                }
                Err(err)
            }
        }
    }

    fn try_apply_policy_direct(policy: UsbPolicy) -> Result<(), UsbControlError> {
        // 1. Aplicar Removable Storage Access Policy
        apply_removable_storage_policy(policy)?;

        // 2. Aplicar USBSTOR como reforco
        apply_usbstor_policy(policy)?;

        // 3. Aplicar WriteProtect se necessario
        if policy == UsbPolicy::Readonly {
            apply_write_protect(true)?;
        } else {
            apply_write_protect(false)?;
        }

        Ok(())
    }

    fn apply_removable_storage_policy(policy: UsbPolicy) -> Result<(), UsbControlError> {
        let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);

        let full_path = format!(r"{}\{}", REMOVABLE_STORAGE_PATH, REMOVABLE_STORAGE_GUID);

        match policy {
            UsbPolicy::Allow => {
                // Tenta remover as restricoes, se existirem
                if let Ok(key) = hklm.open_subkey_with_flags(&full_path, KEY_ALL_ACCESS) {
                    let _ = key.delete_value("Deny_Read");
                    let _ = key.delete_value("Deny_Write");
                    let _ = key.delete_value("Deny_Execute");
                }
                // Tenta remover a chave inteira se estiver vazia
                let _ = hklm.delete_subkey(&full_path);
            }
            UsbPolicy::BlockAll => {
                let (key, _) = hklm
                    .create_subkey(&full_path)
                    .map_err(map_winreg_error)?;

                key.set_value("Deny_Read", &1u32)
                    .map_err(map_winreg_error)?;
                key.set_value("Deny_Write", &1u32)
                    .map_err(map_winreg_error)?;
                key.set_value("Deny_Execute", &1u32)
                    .map_err(map_winreg_error)?;
            }
            UsbPolicy::Readonly => {
                let (key, _) = hklm
                    .create_subkey(&full_path)
                    .map_err(map_winreg_error)?;

                // Permite leitura, bloqueia escrita
                key.set_value("Deny_Read", &0u32)
                    .map_err(map_winreg_error)?;
                key.set_value("Deny_Write", &1u32)
                    .map_err(map_winreg_error)?;
                key.set_value("Deny_Execute", &0u32)
                    .map_err(map_winreg_error)?;
            }
        }

        Ok(())
    }

    fn apply_usbstor_policy(policy: UsbPolicy) -> Result<(), UsbControlError> {
        let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);

        let key = hklm
            .open_subkey_with_flags(USBSTOR_PATH, KEY_ALL_ACCESS)
            .map_err(map_winreg_error)?;

        match policy {
            UsbPolicy::Allow => {
                // Start = 3 habilita o driver
                key.set_value("Start", &3u32)
                    .map_err(map_winreg_error)?;
            }
            UsbPolicy::BlockAll | UsbPolicy::Readonly => {
                // Start = 4 desabilita o driver
                // Nota: Para Readonly, mantemos o driver ativo mas com WriteProtect
                // Porem, como fallback de seguranca, desabilitamos para BlockAll
                if policy == UsbPolicy::BlockAll {
                    key.set_value("Start", &4u32)
                        .map_err(map_winreg_error)?;
                } else {
                    // Readonly mantem driver ativo
                    key.set_value("Start", &3u32)
                        .map_err(map_winreg_error)?;
                }
            }
        }

        Ok(())
    }

    fn apply_write_protect(enable: bool) -> Result<(), UsbControlError> {
        let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);

        if enable {
            let (key, _) = hklm
                .create_subkey(STORAGE_POLICY_PATH)
                .map_err(map_winreg_error)?;

            key.set_value("WriteProtect", &1u32)
                .map_err(map_winreg_error)?;
        } else if let Ok(key) = hklm.open_subkey_with_flags(STORAGE_POLICY_PATH, KEY_ALL_ACCESS) {
            let _ = key.set_value("WriteProtect", &0u32);
        }

        Ok(())
    }

    pub fn get_current_policy() -> Result<UsbPolicy, UsbControlError> {
        let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);

        // Verifica Removable Storage Policy primeiro
        let full_path = format!(r"{}\{}", REMOVABLE_STORAGE_PATH, REMOVABLE_STORAGE_GUID);

        if let Ok(key) = hklm.open_subkey_with_flags(&full_path, KEY_READ) {
            let deny_read: u32 = key.get_value("Deny_Read").unwrap_or(0);
            let deny_write: u32 = key.get_value("Deny_Write").unwrap_or(0);

            if deny_read == 1 && deny_write == 1 {
                return Ok(UsbPolicy::BlockAll);
            }

            if deny_read == 0 && deny_write == 1 {
                return Ok(UsbPolicy::Readonly);
            }
        }

        // Verifica USBSTOR como fallback
        if let Ok(key) = hklm.open_subkey_with_flags(USBSTOR_PATH, KEY_READ) {
            let start: u32 = key.get_value("Start").unwrap_or(3);
            if start == 4 {
                return Ok(UsbPolicy::BlockAll);
            }
        }

        Ok(UsbPolicy::Allow)
    }

    fn is_permission_error(error: &UsbControlError) -> bool {
        match error {
            UsbControlError::PermissionDenied => true,
            UsbControlError::RegistryError(msg) => {
                let lower = msg.to_lowercase();
                lower.contains("access is denied") || lower.contains("acesso negado") || lower.contains("5")
            }
            _ => false,
        }
    }

    #[allow(dead_code)]
    fn apply_policy_with_elevation(policy: UsbPolicy) -> Result<(), UsbControlError> {
        // Cria script temporário para aplicar as chaves via PowerShell elevado
        let temp_dir = std::env::temp_dir();
        let script_path: PathBuf = temp_dir.join("raven_usb_policy.ps1");

        let policy_str = policy.as_str();
        let script = format!(
            r#"$ErrorActionPreference = 'Stop'
$guid = '{guid}'
$policy = '{policy}'

function Set-Allow {{
  reg delete 'HKLM\Software\Policies\Microsoft\Windows\RemovableStorageDevices\{guid}' /f 2>$null
  reg delete 'HKLM\SYSTEM\CurrentControlSet\Control\StorageDevicePolicies' /f 2>$null
  reg add 'HKLM\SYSTEM\CurrentControlSet\Services\USBSTOR' /v Start /t REG_DWORD /d 3 /f | Out-Null
}}

function Set-BlockAll {{
  reg add 'HKLM\Software\Policies\Microsoft\Windows\RemovableStorageDevices\{guid}' /f | Out-Null
  reg add 'HKLM\Software\Policies\Microsoft\Windows\RemovableStorageDevices\{guid}' /v Deny_Read /t REG_DWORD /d 1 /f | Out-Null
  reg add 'HKLM\Software\Policies\Microsoft\Windows\RemovableStorageDevices\{guid}' /v Deny_Write /t REG_DWORD /d 1 /f | Out-Null
  reg add 'HKLM\Software\Policies\Microsoft\Windows\RemovableStorageDevices\{guid}' /v Deny_Execute /t REG_DWORD /d 1 /f | Out-Null
  reg add 'HKLM\SYSTEM\CurrentControlSet\Services\USBSTOR' /v Start /t REG_DWORD /d 4 /f | Out-Null
  reg add 'HKLM\SYSTEM\CurrentControlSet\Control\StorageDevicePolicies' /f | Out-Null
  reg add 'HKLM\SYSTEM\CurrentControlSet\Control\StorageDevicePolicies' /v WriteProtect /t REG_DWORD /d 0 /f | Out-Null
}}

function Set-Readonly {{
  reg add 'HKLM\Software\Policies\Microsoft\Windows\RemovableStorageDevices\{guid}' /f | Out-Null
  reg add 'HKLM\Software\Policies\Microsoft\Windows\RemovableStorageDevices\{guid}' /v Deny_Read /t REG_DWORD /d 0 /f | Out-Null
  reg add 'HKLM\Software\Policies\Microsoft\Windows\RemovableStorageDevices\{guid}' /v Deny_Write /t REG_DWORD /d 1 /f | Out-Null
  reg add 'HKLM\Software\Policies\Microsoft\Windows\RemovableStorageDevices\{guid}' /v Deny_Execute /t REG_DWORD /d 0 /f | Out-Null
  reg add 'HKLM\SYSTEM\CurrentControlSet\Services\USBSTOR' /v Start /t REG_DWORD /d 3 /f | Out-Null
  reg add 'HKLM\SYSTEM\CurrentControlSet\Control\StorageDevicePolicies' /f | Out-Null
  reg add 'HKLM\SYSTEM\CurrentControlSet\Control\StorageDevicePolicies' /v WriteProtect /t REG_DWORD /d 1 /f | Out-Null
}}

switch ($policy) {{
  'ALLOW' {{ Set-Allow }}
  'BLOCK_ALL' {{ Set-BlockAll }}
  'READONLY' {{ Set-Readonly }}
  default {{ throw 'Politica invalida' }}
}}

try {{
  gpupdate /target:computer /force | Out-Null
}} catch {{}}
"#,
            guid = REMOVABLE_STORAGE_GUID,
            policy = policy_str
        );

        fs::write(&script_path, script).map_err(UsbControlError::Io)?;

        // Start-Process com RunAs para acionar UAC
        let arg = format!(
            "Start-Process -WindowStyle Hidden -FilePath powershell -Verb RunAs -Wait -ArgumentList '-ExecutionPolicy Bypass -File \"{}\"'",
            script_path.display()
        );

        let status = Command::new("powershell")
            .arg("-Command")
            .arg(arg)
            .status()
            .map_err(UsbControlError::Io)?;

        if !status.success() {
            return Err(UsbControlError::PermissionDenied);
        }

        Ok(())
    }

    fn map_winreg_error(error: io::Error) -> UsbControlError {
        if let Some(code) = error.raw_os_error() {
            if code == 5 {
                return UsbControlError::PermissionDenied;
            }
        }
        UsbControlError::RegistryError(error.to_string())
    }

    pub fn refresh_group_policy() -> Result<(), UsbControlError> {
        use std::os::windows::process::CommandExt;
        use std::process::Command;

        const CREATE_NO_WINDOW: u32 = 0x08000000;

        // Executa gpupdate para forcar atualizacao das politicas
        let output = Command::new("gpupdate")
            .args(["/target:computer", "/force"])
            .creation_flags(CREATE_NO_WINDOW)
            .output()
            .map_err(UsbControlError::Io)?;

        if !output.status.success() {
            // Nao e critico se falhar, apenas log
            eprintln!(
                "[usb_control] gpupdate retornou erro: {}",
                String::from_utf8_lossy(&output.stderr)
            );
        }

        Ok(())
    }
}

#[cfg(not(target_os = "windows"))]
mod fallback_impl {
    use super::*;

    pub fn apply_usb_policy(_policy: UsbPolicy) -> Result<UsbPolicyResult, UsbControlError> {
        Err(UsbControlError::UnsupportedOs)
    }

    pub fn get_current_policy() -> Result<UsbPolicy, UsbControlError> {
        Err(UsbControlError::UnsupportedOs)
    }

    pub fn refresh_group_policy() -> Result<(), UsbControlError> {
        Err(UsbControlError::UnsupportedOs)
    }
}

#[cfg(target_os = "windows")]
pub use windows_impl::*;

#[cfg(not(target_os = "windows"))]
pub use fallback_impl::*;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_policy_from_str() {
        assert_eq!(UsbPolicy::from_str("ALLOW"), Some(UsbPolicy::Allow));
        assert_eq!(UsbPolicy::from_str("BLOCK_ALL"), Some(UsbPolicy::BlockAll));
        assert_eq!(UsbPolicy::from_str("READONLY"), Some(UsbPolicy::Readonly));
        assert_eq!(UsbPolicy::from_str("allow"), Some(UsbPolicy::Allow));
        assert_eq!(UsbPolicy::from_str("invalid"), None);
    }

    #[test]
    fn test_policy_as_str() {
        assert_eq!(UsbPolicy::Allow.as_str(), "ALLOW");
        assert_eq!(UsbPolicy::BlockAll.as_str(), "BLOCK_ALL");
        assert_eq!(UsbPolicy::Readonly.as_str(), "READONLY");
    }
}