#!/usr/bin/env python3
# Hermes Node Protocol
# Copyright (c) 2026 Stefy (nextime) Lanza <stefy@nexlab.net>
# All rights reserved.
# 
# This software is released under the MIT License with a copyleft clause.
# See the LICENSE file for full terms.
"""
Hermes Node Agent - Reverse-connection node executor

Connects to Hermes Gateway via WebSocket and executes commands.
Supports optional tools: browser control, computer_control.

Author: Lisa (Hermes AI)
Date: 2026-04-30 (enhanced)
"""

import asyncio
import json
import logging
import os
import subprocess
import sys
import time
import argparse
from pathlib import Path
from typing import Optional, Dict, Any, List

try:
    import websockets
except ImportError:
    print("ERROR: websockets library not found. Install with: pip install websockets")
    sys.exit(1)

try:
    from browser_controller import BrowserController
    HAS_BROWSER = True
except ImportError:
    HAS_BROWSER = False

logger = logging.getLogger(__name__)

# ═════════════════════════════════════════════════════════════════════════
# DEFAULT CONFIGURATION
# ═════════════════════════════════════════════════════════════════════════

DEFAULT_GATEWAY_TOKEN = 'GATEWAY_TOKEN_MUST_BE_PROVIDED'

DEFAULT_CONFIG = {
    'gateway_url': 'ws://127.0.0.1:8765',
    'node_name': 'unknown',
    'token': DEFAULT_GATEWAY_TOKEN,
    'sexec_path': '~/.config/hermes-node/sexec/sexec.sh',
    'reconnect_interval': 5,
    'heartbeat_interval': 30,
    'gateway_cert_path': None,
    'capabilities': ['exec'],
    'enable_browser': False,
    'enable_computer_control': False,
}

# ═════════════════════════════════════════════════════════════════════════
# PLATFORM ABSTRACTION LAYER
# ═════════════════════════════════════════════════════════════════════════

class PlatformError(RuntimeError):
    """Raised when platform-specific operations fail."""

def is_windows() -> bool:
    return sys.platform in ('win32', 'cygwin')

def is_linux() -> bool:
    return sys.platform.startswith('linux')

def is_macos() -> bool:
    return sys.platform == 'darwin'

# ═════════════════════════════════════════════════════════════════════════
# COMMAND EXECUTION ABSTRACTION
# ═════════════════════════════════════════════════════════════════════════

class CommandExecutor:
    """Abstract base class for executing commands with permission enforcement."""

    def __init__(self, permission_rules: Dict[str, List[str]]):
        self.permissions = permission_rules or {'allow': [], 'deny': [], 'ask': []}

    def execute(self, command: str, approved: bool = False) -> Dict[str, Any]:
        """Execute command respecting permission rules."""
        raise NotImplementedError

    def _check_permission(self, command: str, approved: bool) -> tuple[bool, str]:
        """Check allow/deny/ask rules.
        Returns (allowed, reason). 'ask' means requires approval gate.
        """
        import re
        cmd = command.strip()
        # Deny (highest priority)
        for pattern in self.permissions.get('deny', []):
            if re.search(pattern, cmd, re.IGNORECASE):
                return False, f"Denied by pattern '{pattern}'"
        # Ask (medium) — only blocks if not approved
        if not approved:
            for pattern in self.permissions.get('ask', []):
                if re.search(pattern, cmd, re.IGNORECASE):
                    return True, 'ask'
        # Allow (explicit)
        if self.permissions.get('allow'):
            for pattern in self.permissions['allow']:
                if re.search(pattern, cmd, re.IGNORECASE):
                    return True, 'allowed'
            return False, "Not in allow list"
        return True, 'default-allow'


# ── POSIX ──────────────────────────────────────────────────────────────────
class PosixCommandExecutor(CommandExecutor):
    """POSIX implementation — uses sexec.sh and environment variable."""

    def __init__(self, sexec_path: str, permission_rules: Dict[str, List[str]]):
        super().__init__(permission_rules)
        self.sexec_path = Path(sexec_path).expanduser()

    def execute(self, command: str, approved: bool = False) -> Dict[str, Any]:
        allowed, reason = self._check_permission(command, approved)
        if not allowed and reason != 'ask':
            return {'success': False, 'error': f'Permission denied: {reason}', 'exit_code': 127}
        if not approved and reason == 'ask':
            return {'success': False, 'error': 'Command requires approval', 'exit_code': 2}
        if not self.sexec_path.exists():
            return {'success': False, 'error': f'sexec not found: {self.sexec_path}', 'exit_code': 127}
        try:
            env = os.environ.copy()
            env['SEXEC_COMMAND'] = command
            proc = subprocess.Popen(
                [str(self.sexec_path)],
                stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env, text=True
            )
            out, err = proc.communicate(timeout=300)
            return {'success': proc.returncode == 0, 'stdout': out, 'stderr': err,
                    'exit_code': proc.returncode}
        except Exception as e:
            return {'success': False, 'error': str(e), 'exit_code': -1}


# ── WINDOWS ────────────────────────────────────────────────────────────────
class WindowsCommandExecutor(CommandExecutor):
    """Windows implementation using PowerShell with base64-encoded commands.

    Design:
    - Uses PowerShell 7+ (pwsh.exe) if available, else Windows PowerShell
    - Encodes command as base64(utf-16le) to avoid shell quoting issues
    - No-Persist policy flags make it a child process
    - Returns stdout/stderr as text with exit code
    """

    def __init__(self, permission_rules: Dict[str, List[str]]):
        super().__init__(permission_rules)
        self.powershell = self._find_powershell()

    def _find_powershell(self) -> str:
        for candidate in [
            r'C:\Program Files\PowerShell\7\pwsh.exe',
            r'C:\Program Files (x86)\PowerShell\7\pwsh.exe',
            r'C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe',
        ]:
            if Path(candidate).exists():
                return candidate
        return 'powershell'

    def execute(self, command: str, approved: bool = False) -> Dict[str, Any]:
        allowed, reason = self._check_permission(command, approved)
        if not allowed and reason != 'ask':
            return {'success': False, 'error': f'Permission denied: {reason}', 'exit_code': 127}
        if not approved and reason == 'ask':
            return {'success': False, 'error': 'Command requires approval', 'exit_code': 2}
        try:
            import base64
            encoded = base64.b64encode(command.encode('utf-16le')).decode('ascii')
            proc = subprocess.Popen(
                [self.powershell, '-NoProfile', '-NonInteractive',
                 '-ExecutionPolicy', 'Bypass', '-EncodedCommand', encoded],
                stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
                creationflags=subprocess.CREATE_NO_WINDOW if hasattr(subprocess, 'CREATE_NO_WINDOW') else 0
            )
            out, err = proc.communicate(timeout=300)
            return {'success': proc.returncode == 0, 'stdout': out, 'stderr': err,
                    'exit_code': proc.returncode}
        except FileNotFoundError:
            return {'success': False, 'error': 'PowerShell not found', 'exit_code': 127}
        except Exception as e:
            return {'success': False, 'error': str(e), 'exit_code': -1}


# ═════════════════════════════════════════════════════════════════════════
# COMPUTER CONTROL LAYER — cross-platform abstraction
# ═════════════════════════════════════════════════════════════════════════

class ComputerControllerBase:
    """Base class for desktop automation, platform-agnostic API."""

    def screenshot(self, output_path: Optional[str] = None) -> Dict[str, Any]:
        raise NotImplementedError

    def mouse_move(self, x: int, y: int) -> Dict[str, Any]:
        raise NotImplementedError

    def mouse_click(self, button: int = 1) -> Dict[str, Any]:
        raise NotImplementedError

    def mouse_position(self) -> Dict[str, Any]:
        raise NotImplementedError

    def type_text(self, text: str) -> Dict[str, Any]:
        raise NotImplementedError

    def key_press(self, key: str) -> Dict[str, Any]:
        raise NotImplementedError

    def get_active_window(self) -> Dict[str, Any]:
        raise NotImplementedError


# ── POSIX computer control (xdotool + import/ImageMagick) ──────────────────
class PosixComputerController(ComputerControllerBase):
    """Linux X11 automation via command-line tools."""

    def __init__(self):
        self.display = os.environ.get('DISPLAY', ':0')

    def screenshot(self, output_path: Optional[str] = None) -> Dict[str, Any]:
        if output_path:
            r = self._run(f'import -display {self.display} -window root "{output_path}"')
            return {'success': r['success'], 'path': output_path, 'error': r.get('error')}
        else:
            import base64
            try:
                result = subprocess.run(
                    f'import -display {self.display} -window root png:-',
                    shell=True, capture_output=True, timeout=30
                )
                if result.returncode == 0:
                    return {
                        'success': True, 'format': 'png',
                        'data': base64.b64encode(result.stdout).decode('ascii'),
                        'size': len(result.stdout)
                    }
            except Exception as e:
                return {'success': False, 'error': str(e)}
            return {'success': False, 'error': 'screenshot failed'}

    def mouse_move(self, x: int, y: int) -> Dict[str, Any]:
        return self._run(f'xdotool mousemove {x} {y}')

    def mouse_click(self, button: int = 1) -> Dict[str, Any]:
        return self._run(f'xdotool click {button}')

    def mouse_position(self) -> Dict[str, Any]:
        out = self._run('xdotool getmouselocation --shell')
        pos = {}
        if out['success']:
            for line in out['stdout'].splitlines():
                if '=' in line:
                    k, v = line.split('=', 1)
                    pos[k] = int(v)
        return {'success': out['success'], 'position': pos, 'error': out.get('error')}

    def type_text(self, text: str) -> Dict[str, Any]:
        # Escape single quotes for shell
        safe = text.replace("'", "'\"'\"'")
        return self._run(f"xdotool type --delay 1 '{safe}'")

    def key_press(self, key: str) -> Dict[str, Any]:
        return self._run(f'xdotool key {key}')

    def get_active_window(self) -> Dict[str, Any]:
        win_id = self._run('xdotool getactivewindow')
        if win_id['success']:
            title = self._run(f'xdotool getwindowname {win_id["stdout"]}')
            return {'success': True, 'window_id': win_id['stdout'],
                    'title': title.get('stdout', ''), 'error': title.get('error')}
        return win_id

    def _run(self, cmd: str) -> Dict[str, Any]:
        try:
            r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
            return {'success': r.returncode == 0, 'stdout': r.stdout.strip(),
                    'stderr': r.stderr.strip(), 'exit_code': r.returncode}
        except Exception as e:
            return {'success': False, 'error': str(e)}


# ── WINDOWS computer control (pyautogui + PIL) ────────────────────────────
class WindowsComputerController(ComputerControllerBase):
    """Windows desktop automation using pyautogui and PIL/Pillow."""

    def __init__(self):
        try:
            import PIL.ImageGrab
            import pyautogui
            self.ImageGrab = PIL.ImageGrab
            self.pyautogui = pyautogui
            pyautogui.FAILSAFE = False  # Allow user to abort by moving mouse to corner
        except ImportError as e:
            raise ImportError(f"Windows computer_control requires pyautogui and Pillow: {e}")

    def screenshot(self, output_path: Optional[str] = None) -> Dict[str, Any]:
        try:
            img = self.ImageGrab.grab()
            if output_path:
                img.save(output_path)
                return {'success': True, 'path': output_path}
            import io, base64
            buf = io.BytesIO()
            img.save(buf, format='PNG')
            return {
                'success': True, 'format': 'png',
                'data': base64.b64encode(buf.getvalue()).decode('ascii'),
                'size': len(buf.getvalue())
            }
        except Exception as e:
            return {'success': False, 'error': str(e)}

    def mouse_move(self, x: int, y: int) -> Dict[str, Any]:
        try:
            self.pyautogui.moveTo(x, y)
            return {'success': True}
        except Exception as e:
            return {'success': False, 'error': str(e)}

    def mouse_click(self, button: int = 1) -> Dict[str, Any]:
        try:
            btn = {1: 'left', 2: 'middle', 3: 'right'}.get(button, 'left')
            self.pyautogui.click(button=btn)
            return {'success': True}
        except Exception as e:
            return {'success': False, 'error': str(e)}

    def mouse_position(self) -> Dict[str, Any]:
        try:
            x, y = self.pyautogui.position()
            return {'success': True, 'position': {'x': x, 'y': y}}
        except Exception as e:
            return {'success': False, 'error': str(e)}

    def type_text(self, text: str) -> Dict[str, Any]:
        try:
            self.pyautogui.write(text, interval=0.01)
            return {'success': True}
        except Exception as e:
            return {'success': False, 'error': str(e)}

    def key_press(self, key: str) -> Dict[str, Any]:
        try:
            key_map = {
                'return': 'enter', 'enter': 'enter', 'esc': 'escape',
                'ctrl': 'ctrl', 'alt': 'alt', 'shift': 'shift',
                'tab': 'tab', 'space': 'space', 'backspace': 'backspace',
                'delete': 'delete', 'up': 'up', 'down': 'down',
                'left': 'left', 'right': 'right', 'home': 'home', 'end': 'end'
            }
            mapped = key_map.get(key.lower(), key)
            self.pyautogui.press(mapped)
            return {'success': True}
        except Exception as e:
            return {'success': False, 'error': str(e)}

    def get_active_window(self) -> Dict[str, Any]:
        try:
            import win32gui, win32process
            hwnd = win32gui.GetForegroundWindow()
            title = win32gui.GetWindowText(hwnd)
            _, pid = win32process.GetWindowThreadProcessId(hwnd)
            return {
                'success': True,
                'window_id': hwnd,
                'title': title,
                'process_id': pid
            }
        except ImportError:
            return {'success': True, 'title': self.pyautogui.getActiveWindow().title}
        except Exception as e:
            return {'success': False, 'error': str(e)}


# ── Factory functions ─────────────────────────────────────────────────────
def make_executor(config: Dict[str, Any]) -> CommandExecutor:
    """Select appropriate command executor for current platform."""
    perms = config.get('permissions', {})
    if is_windows():
        return WindowsCommandExecutor(perms)
    else:
        sexec = config.get('sexec_path', str(Path.home() / '.openclaw/skills/sexec/sexec.sh'))
        return PosixCommandExecutor(sexec, perms)

def make_computer_controller(config: Dict[str, Any]) -> Optional[ComputerControllerBase]:
    """Select and instantiate the appropriate computer controller, or None if deps missing."""
    if not config.get('enable_computer_control'):
        return None
    if is_windows():
        try:
            return WindowsComputerController()
        except ImportError as e:
            logger.warning(f"computer_control disabled (missing deps): {e}")
            return None
    else:
        # Linux/macOS
        if is_macos():
            logger.warning("macOS computer_control not implemented yet")
            return None
        # Linux
        if subprocess.run(['which', 'xdotool'], capture_output=True).returncode != 0:
            logger.warning("xdotool not found — computer_control disabled")
            return None
        try:
            return PosixComputerController()
        except Exception as e:
            logger.warning(f"computer_control init failed: {e}")
            return None


class NodeAgent:
    def __init__(self, config_path: Optional[str] = None):
        self.config = self._load_config(config_path)
        self.capabilities = self._detect_capabilities()
        self.computer = None
        if self.capabilities.get('enable_computer_control'):
            self.computer = ComputerController()
        self.browser = None
        if self.capabilities.get('enable_browser') and HAS_BROWSER:
            try:
                self.browser = BrowserController()
            except Exception as e:
                logger.warning(f"BrowserController init failed: {e}")

    def _load_config(self, path: Optional[str]) -> Dict[str, Any]:
        """Load node configuration from JSON."""
        cfg_path = Path(path).expanduser() if path else Path.home() / '.config' / 'hermes-node' / 'config.json'
        if not cfg_path.exists():
            logger.error(f"Config not found: {cfg_path}")
            sys.exit(1)
        with open(cfg_path) as f:
            data = json.load(f)
        # Merge defaults
        merged = {**DEFAULT_CONFIG, **data}
        if not merged['token']:
            logger.error("Token missing from config")
            sys.exit(1)
        return merged

    def _detect_capabilities(self) -> Dict[str, Any]:
        """Detect which optional tools are available on this machine."""
        caps = {
            'enable_browser': self.config.get('enable_browser', False),
            'enable_computer_control': self.config.get('enable_computer_control', False),
        }
        # Browser detection: extension installed signals availability via separate channel
        # Computer control: check system tools
        if caps['enable_computer_control']:
            cc_info = {
                'display': os.environ.get('DISPLAY', ':0'),
                'has_xdotool': subprocess.run(['which', 'xdotool'], capture_output=True).returncode == 0,
                'has_import': subprocess.run(['which', 'import'], capture_output=True).returncode == 0,
                'has_scrot': subprocess.run(['which', 'scrot'], capture_output=True).returncode == 0,
            }
            caps['computer_control'] = cc_info
        return caps

    async def connect_and_run(self):
        """Main loop: connect to gateway and process commands."""
        url = f"{self.config['gateway_url']}?node_name={self.config['node_name']}&token={self.config['token']}"
        logger.info(f"Connecting to {url}")

        while True:
            try:
                async with websockets.connect(url, ping_interval=20, ping_timeout=10) as ws:
                    logger.info("Connected to gateway — awaiting commands")

                    # Send capabilities announcement
                    await ws.send(json.dumps({
                        "type": "capabilities",
                        "node_name": self.config['node_name'],
                        "tools": self._get_available_tools(),
                        "capabilities": self.capabilities
                    }))

                    # Listen for commands
                    async for raw in ws:
                        msg = json.loads(raw)
                        await self._handle_message(ws, msg)

            except Exception as e:
                logger.error(f"Connection error: {e}")
                logger.info(f"Reconnecting in {self.config['reconnect_interval']}s...")
                await asyncio.sleep(self.config['reconnect_interval'])

    def _get_available_tools(self) -> list:
        """Return list of capability strings for gateway registration."""
        tools = ['exec']  # always present
        if self.capabilities.get('enable_browser'):
            tools.append('browser_control')
        if self.capabilities.get('enable_computer_control'):
            tools.append('computer_control')
        return tools

    async def _handle_message(self, ws, msg: Dict[str, Any]):
        req_type = msg.get('type')
        if req_type == 'exec':
            cmd_id = msg['id']
            command = msg['command']
            await self._handle_exec(ws, cmd_id, command)
        elif req_type == 'computer_control':
            action = msg['action']
            params = msg.get('params', {})
            await self._handle_cc(ws, msg.get('id'), action, params)
        else:
            logger.warning(f"Unknown message type: {req_type}")

    async def _handle_exec(self, ws, cmd_id: str, command: str):
        """Execute a shell command via sexec (respecting permissions)."""
        logger.info(f"Exec: {command}")
        try:
            sexec = Path(self.config['sexec_path']).expanduser()
            if not sexec.exists():
                await ws.send(json.dumps({
                    "type": "exec_result",
                    "id": cmd_id,
                    "error": f"sexec not found at {sexec}",
                    "stdout": "", "stderr": "", "exit_code": 127
                }))
                return

            env = os.environ.copy()
            env['SEXEC_COMMAND'] = command
            proc = subprocess.Popen(
                [sexec],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                env=env,
                text=True
            )
            stdout, stderr = proc.communicate(timeout=300)
            await ws.send(json.dumps({
                "type": "exec_result",
                "id": cmd_id,
                "command": command,
                "stdout": stdout,
                "stderr": stderr,
                "exit_code": proc.returncode
            }))
        except subprocess.CalledProcessError as e:
            await ws.send(json.dumps({
                "type": "exec_result",
                "id": cmd_id,
                "error": str(e),
                "exit_code": -1
            }))


def main():
    parser = argparse.ArgumentParser(description="Hermes Node Agent")
    parser.add_argument('--config', type=str, help='Path to config JSON')
    parser.add_argument('--debug', action='store_true', help='Debug logging')
    args = parser.parse_args()

    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    # Load config to check token
    config = NodeAgent(args.config)._load_config(args.config)
    if config['token'] == DEFAULT_GATEWAY_TOKEN or config['token'] == 'GATEWAY_TOKEN_MUST_BE_PROVIDED':
        logger.error("ERROR: Token not set in config. Edit ~/.config/hermes-node/config.json")
        sys.exit(1)

    logger.info(f"Node '{config['node_name']}' starting — tools: {NodeAgent(args.config)._get_available_tools()}")

    agent = NodeAgent(args.config)
    try:
        asyncio.run(agent.connect_and_run())
    except KeyboardInterrupt:
        logger.info("Shutting down")

if __name__ == '__main__':
    main()
