#!/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 ssl
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__)

LOG_PREVIEW_LEN = 120


def _preview_command(command: Any, limit: int = LOG_PREVIEW_LEN) -> str:
    """Return a compact single-line preview for logging."""
    if isinstance(command, (list, tuple)):
        text = ' '.join(str(part) for part in command)
    else:
        text = str(command)
    text = ' '.join(text.split())
    if len(text) > limit:
        return text[:limit] + '…'
    return text


def _setup_logging(debug: bool = False) -> None:
    level = logging.DEBUG if debug else logging.INFO
    logging.basicConfig(level=level, format='%(message)s')


def _log_start(node_name: str, tools: list) -> None:
    logger.info(f"start ▶ {node_name} — {', '.join(tools)}")


def _log_connect(url: str) -> None:
    logger.info(f"connect ▶ {url}")


def _log_tls_disabled() -> None:
    logger.info("tls verify disabled")


def _log_connected() -> None:
    logger.info("connect ✓")


def _log_registering(node_name: str) -> None:
    logger.info(f"register ▶ {node_name}")


def _log_registered(node_name: str) -> None:
    logger.info(f"register ✓ {node_name}")


def _log_waiting() -> None:
    logger.info("waiting for commands")


def _log_exec_received(command: Any) -> None:
    logger.info(f"exec ▶ {_preview_command(command)}")


def _log_exec_completed(command: Any, exit_code: Any) -> None:
    logger.info(f"exec ✓ exit={exit_code} — {_preview_command(command)}")


def _log_exec_failed(command: Any, error: Any, exit_code: Any = None) -> None:
    prefix = f"exec ✗ exit={exit_code}" if exit_code is not None else "exec ✗"
    logger.error(f"{prefix} — {_preview_command(command)} — {error}")


def _log_browser_received(command: Any) -> None:
    logger.info(f"browser ▶ {_preview_command(command)}")


def _log_cc_received(action: Any) -> None:
    logger.info(f"computer ▶ {_preview_command(action)}")


def _log_reconnect(delay: Any, reason: Any) -> None:
    logger.info(f"reconnect in {delay}s — {reason}")


def _log_registration_ack() -> None:
    logger.debug("register ack")


def _log_heartbeat_ack() -> None:
    logger.debug("heartbeat ack")


def _log_unknown_message(req_type: Any) -> None:
    logger.warning(f"unknown message: {req_type}")


def _log_connection_error(message: Any) -> None:
    logger.error(f"connection error — {message}")


def _log_config_missing(path: Path) -> None:
    logger.error(f"config missing — {path}")


def _log_token_missing() -> None:
    logger.error("token missing in config")


def _log_init_warning(component: str, message: Any) -> None:
    logger.warning(f"{component} init failed — {message}")


def _log_disabled(component: str, message: str) -> None:
    logger.warning(f"{component} disabled — {message}")


def _log_shutdown() -> None:
    logger.info("shutdown")

# ═════════════════════════════════════════════════════════════════════════
# 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: Any, approved: bool) -> tuple[bool, str]:
        """Check allow/deny/ask rules.
        Returns (allowed, reason). 'ask' means requires approval gate.
        Accepts command as string or argv list.
        """
        import re
        if isinstance(command, (list, tuple)):
            cmd = ' '.join(str(part) for part in command)
        else:
            cmd = str(command)
        cmd = cmd.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 using integrated permission checks and /bin/sh."""

    def execute(self, command: Any, 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 isinstance(command, (list, tuple)):
            cmd = ' '.join(subprocess.list2cmdline([str(part)]) if is_windows() else __import__('shlex').quote(str(part)) for part in command)
        else:
            cmd = str(command)
        try:
            proc = subprocess.Popen(
                ['/bin/sh', '-c', cmd],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            out, err = proc.communicate(timeout=300)
            return {
                'success': proc.returncode == 0,
                'stdout': out,
                'stderr': err,
                'exit_code': proc.returncode,
            }
        except subprocess.TimeoutExpired:
            try:
                proc.kill()
            except Exception:
                pass
            return {'success': False, 'error': 'Command timed out', 'exit_code': 124}
        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:
        return PosixCommandExecutor(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:
            _log_disabled('computer_control', f"missing deps: {e}")
            return None
    else:
        # Linux/macOS
        if is_macos():
            _log_disabled('computer_control', 'macOS not implemented yet')
            return None
        # Linux
        if subprocess.run(['which', 'xdotool'], capture_output=True).returncode != 0:
            _log_disabled('computer_control', 'xdotool not found')
            return None
        try:
            return PosixComputerController()
        except Exception as e:
            _log_init_warning('computer_control', e)
            return None


class NodeAgent:
    def __init__(self, config_path: Optional[str] = None):
        self.config = self._load_config(config_path)
        self.executor = make_executor(self.config)
        self.capabilities = self._detect_capabilities()
        self.computer = make_computer_controller(self.config)
        self.browser = None
        if self.capabilities.get('enable_browser') and HAS_BROWSER:
            try:
                self.browser = BrowserController()
            except Exception as e:
                _log_init_warning('browser_control', 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():
            _log_config_missing(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']:
            _log_token_missing()
            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']}"
        _log_connect(url)

        ssl_context = None
        if self.config['gateway_url'].startswith('wss://'):
            ssl_context = ssl.create_default_context()
            ssl_context.check_hostname = False
            ssl_context.verify_mode = ssl.CERT_NONE
            _log_tls_disabled()

        while True:
            try:
                async with websockets.connect(url, ping_interval=20, ping_timeout=10, ssl=ssl_context) as ws:
                    _log_connected()

                    # Send registration expected by gateway
                    _log_registering(self.config['node_name'])
                    await ws.send(json.dumps({
                        "type": "register",
                        "node_name": self.config['node_name'],
                        "version": "node-agent-clean",
                        "tools": self._get_available_tools(),
                        "capabilities": self.capabilities
                    }))
                    _log_waiting()

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

            except Exception as e:
                _log_connection_error(e)
                _log_reconnect(self.config['reconnect_interval'], e)
                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)
        elif req_type == 'register_ack':
            _log_registered(self.config['node_name'])
            _log_registration_ack()
        elif req_type == 'heartbeat_ack':
            _log_heartbeat_ack()
        else:
            _log_unknown_message(req_type)

    async def _handle_exec(self, ws, cmd_id: str, command: str):
        """Execute a shell command and reply with the gateway's exec protocol."""
        _log_exec_received(command)
        result = self.executor.execute(command)
        stdout = result.get('stdout', '') or ''
        stderr = result.get('stderr', '') or ''
        if stdout:
            await ws.send(json.dumps({
                "type": "exec_output",
                "id": cmd_id,
                "stream": "stdout",
                "data": stdout
            }))
        if stderr:
            await ws.send(json.dumps({
                "type": "exec_output",
                "id": cmd_id,
                "stream": "stderr",
                "data": stderr
            }))
        exit_code = result.get('exit_code', -1)
        error = result.get('error')
        if error:
            _log_exec_failed(command, error, exit_code)
        else:
            _log_exec_completed(command, exit_code)
        await ws.send(json.dumps({
            "type": "exec_complete",
            "id": cmd_id,
            "exit_code": exit_code,
            "error": error
        }))


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

    _setup_logging(args.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':
        _log_token_missing()
        sys.exit(1)

    tools = NodeAgent(args.config)._get_available_tools()
    _log_start(config['node_name'], tools)

    agent = NodeAgent(args.config)
    try:
        asyncio.run(agent.connect_and_run())
    except KeyboardInterrupt:
        _log_shutdown()

if __name__ == '__main__':
    main()
