//! Modulo USB Policy - Controle de dispositivos USB
//!
//! Implementa o controle de armazenamento USB no Windows.
//! Como o servico roda como LocalSystem, nao precisa de elevacao.

use serde::{Deserialize, Serialize};
use std::io;
use thiserror::Error;
use tracing::{error, info, warn};
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";

#[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)]
pub enum UsbControlError {
    #[error("Politica USB invalida: {0}")]
    InvalidPolicy(String),

    #[error("Erro de registro do Windows: {0}")]
    RegistryError(String),

    #[error("Permissao negada")]
    PermissionDenied,

    #[error("Erro de I/O: {0}")]
    Io(#[from] io::Error),
}

/// Aplica uma politica de USB
pub fn apply_policy(policy_str: &str) -> Result<UsbPolicyResult, UsbControlError> {
    let policy = UsbPolicy::from_str(policy_str)
        .ok_or_else(|| UsbControlError::InvalidPolicy(policy_str.to_string()))?;

    let now = chrono::Utc::now().timestamp_millis();

    info!("Aplicando politica USB: {:?}", policy);

    // 1. Aplicar Removable Storage Policy
    apply_removable_storage_policy(policy)?;

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

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

    // 4. Atualizar Group Policy (opcional)
    if let Err(e) = refresh_group_policy() {
        warn!("Falha ao atualizar group policy: {}", e);
    }

    info!("Politica USB aplicada com sucesso: {:?}", policy);

    Ok(UsbPolicyResult {
        success: true,
        policy: policy.as_str().to_string(),
        error: None,
        applied_at: Some(now),
    })
}

/// Retorna a politica USB atual
pub fn get_current_policy() -> Result<String, 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("BLOCK_ALL".to_string());
        }

        if deny_read == 0 && deny_write == 1 {
            return Ok("READONLY".to_string());
        }
    }

    // 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("BLOCK_ALL".to_string());
        }
    }

    Ok("ALLOW".to_string())
}

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 => {
            // Start = 4 desabilita o driver
            key.set_value("Start", &4u32)
                .map_err(map_winreg_error)?;
        }
        UsbPolicy::Readonly => {
            // 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(())
}

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

    const CREATE_NO_WINDOW: u32 = 0x08000000;

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

    if !output.status.success() {
        warn!(
            "gpupdate retornou erro: {}",
            String::from_utf8_lossy(&output.stderr)
        );
    }

    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())
}