import os
import subprocess

async def screenshot_tool(*args, **kwargs):
    if os.environ.get('CLOUD_ENV', 'false').lower() == 'true':
        return "Warning: Cannot take screenshot from Cloud Environment. Please run JARVIS locally for this feature."
    try:
        from PIL import ImageGrab
        img = ImageGrab.grab()
        path = os.path.join(os.path.dirname(__file__), '..', '..', 'scratch', 'screenshot.png')
        img.save(path)
        return f"Screenshot saved to {path}"
    except Exception as e:
        return f"Failed to take screenshot: {e}"

async def open_app_tool(app_name: str = "", *args, **kwargs):
    if os.environ.get('CLOUD_ENV', 'false').lower() == 'true':
        from backend.ws.agent_ws import ws_manager
        await ws_manager.broadcast({
            "event": "system:execute",
            "payload": {"cmd": f"open_app::{app_name}"}
        })
        return f"Dispatched request to open {app_name} to connected PC Relay Client."
    try:
        subprocess.Popen(f"start {app_name}", shell=True)
        return f"Successfully opened {app_name} on local PC."
    except Exception as e:
        return f"Failed to open app: {e}"

async def usb_devices_tool(*args, **kwargs):
    from backend.services.usb_monitor import get_active_usb_drives
    if os.environ.get('CLOUD_ENV', 'false').lower() == 'true':
        return "Warning: Cannot scan PC USBs from Cloud Environment natively."
    drives = get_active_usb_drives()
    return f"Active USB Drives: {drives}"

async def system_stats_tool(*args, **kwargs):
    import psutil
    if os.environ.get('CLOUD_ENV', 'false').lower() == 'true':
        cpu = psutil.cpu_percent()
        mem = psutil.virtual_memory().percent
        return f"Cloud Server Stats - CPU: {cpu}%, Memory: {mem}%"
    cpu = psutil.cpu_percent()
    mem = psutil.virtual_memory().percent
    return f"Local PC Stats - CPU: {cpu}%, Memory: {mem}%"

async def notification_tool(message: str = "", *args, **kwargs):
    if os.environ.get('CLOUD_ENV', 'false').lower() == 'true':
        from backend.ws.agent_ws import ws_manager
        await ws_manager.broadcast({
            "event": "system:notify",
            "payload": {"message": message, "speak": False}
        })
        return f"Dispatched notification '{message}' to connected PC Relay Client."
    try:
        cmd = f'''powershell -Command "Add-Type -AssemblyName System.Windows.Forms; [System.Windows.Forms.MessageBox]::Show('{message}')"'''
        subprocess.Popen(cmd, shell=True)
        return f"Displayed notification: {message}"
    except Exception as e:
        return f"Failed to show notification: {e}"

async def backup_vault_tool(drive_letter: str = None, *args, **kwargs):
    """
    Backs up the system to the specified USB drive letter.
    If drive_letter is not provided, it will auto-detect the first connected USB.
    """
    import asyncio
    import shutil
    import tempfile
    import os
    from backend.services.usb_vault import encrypt_directory
    from backend.services.usb_monitor import get_active_usb_drives
    
    if not drive_letter:
        active = get_active_usb_drives()
        if not active:
            return "Error: No active USB drives detected."
        drive_letter = active[0].get("drive_letter", "E:")
        
    if not drive_letter.endswith(":"):
        drive_letter += ":"
        
    source_code_path = 'C:\\Users\\Admin\\Desktop\\F.R.I.D.A.Y - OMEGA'
    appdata_path = os.path.join(os.environ.get('APPDATA', os.path.expanduser('~')), 'JARVIS_OS')
    target = f'{drive_letter}\\OMEGA_CORE_V15.vault'
    
    def perform_vault_backup():
        with tempfile.TemporaryDirectory() as staging_dir:
            # 1. Copy the entire codebase
            code_staging = os.path.join(staging_dir, "F.R.I.D.A.Y - OMEGA")
            shutil.copytree(source_code_path, code_staging, dirs_exist_ok=True)
            
            # 2. Copy the AppData database/keys if they exist (EXE check data)
            if os.path.exists(appdata_path):
                appdata_staging = os.path.join(staging_dir, "JARVIS_OS_APPDATA")
                shutil.copytree(appdata_path, appdata_staging, dirs_exist_ok=True)
                
            # 3. Zip and encrypt everything together
            encrypt_directory(staging_dir, target)
            
    await asyncio.to_thread(perform_vault_backup)
    return f"Master Vault Backup ({target}) completed successfully! All code, cloud servers, and EXE databases encrypted."