"""
TokioAI Engine v3.0 -- Server-Side Tool Executor
Executes tools LOCALLY on the Pi without CLI round-trips.
This is what makes the Engine think like Opus: it can investigate autonomously.

Safety: only execute_local, read_file, search_files, diagnose are allowed server-side.
Write operations (write_file, edit_file) are ALWAYS delegated to the CLI for user approval.
Remote operations (execute_raspi, execute_gcp, execute_router) are also delegated.
"""

import os
import re
import subprocess
import time
import glob as glob_module
from typing import Dict, Tuple, Optional


# Tools that the Engine can execute directly (read-only / diagnostic)
SAFE_TOOLS = {
    "execute_local",
    "read_file",
    "search_files",
    "diagnose",
    "robot_vision",
}

# Tools that must be delegated to the CLI (writes, remote)
DELEGATE_TOOLS = {
    "write_file",
    "edit_file",
    "execute_raspi",
    "execute_gcp",
    "execute_router",
    "memory",
    "task",
}

# Commands that are NEVER allowed server-side (destructive)
BLOCKED_COMMANDS = [
    r'\brm\s+-rf\s+/',
    r'\bmkfs\b',
    r'\bdd\s+if=',
    r'\b(reboot|shutdown|poweroff|halt)\b',
    r'\bsudo\s+(rm|mkfs|dd|reboot|shutdown|poweroff|halt)\b',
    r'>\s*/dev/',
    r'\bchmod\s+-R\s+777\s+/',
    r'\bchown\s+-R.*\s+/',
]


def is_safe_tool(tool_name: str) -> bool:
    """Check if a tool can be executed server-side."""
    return tool_name in SAFE_TOOLS


def is_command_blocked(command: str) -> bool:
    """Check if a shell command is blocked for safety."""
    for pattern in BLOCKED_COMMANDS:
        if re.search(pattern, command, re.IGNORECASE):
            return True
    return False


def execute_tool(tool_name: str, arguments: Dict, timeout: int = 30) -> Tuple[bool, str]:
    """
    Execute a tool server-side.
    
    Returns:
        (success: bool, output: str)
        
    If the tool is not in SAFE_TOOLS, returns (False, "DELEGATE") to signal
    that this tool should be sent back to the CLI.
    """
    if tool_name not in SAFE_TOOLS:
        return False, "DELEGATE"
    
    try:
        if tool_name == "execute_local":
            return _execute_local(arguments, timeout)
        elif tool_name == "read_file":
            return _read_file(arguments)
        elif tool_name == "search_files":
            return _search_files(arguments)
        elif tool_name == "diagnose":
            return _diagnose(arguments)
        elif tool_name == "robot_vision":
            return _robot_vision(arguments)
        else:
            return False, f"Unknown tool: {tool_name}"
    except Exception as e:
        return False, f"Error executing {tool_name}: {str(e)}"


def _execute_local(arguments: Dict, default_timeout: int = 30) -> Tuple[bool, str]:
    """Execute a shell command locally."""
    command = arguments.get("command", "")
    if not command:
        return False, "No command specified"
    
    if is_command_blocked(command):
        return False, f"Command blocked for safety: {command}"
    
    timeout = min(arguments.get("timeout", default_timeout), 60)  # Cap at 60s for server-side
    
    try:
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=timeout,
            env={**os.environ, "TERM": "dumb"},  # prevent color codes
        )
        
        output = result.stdout
        if result.stderr:
            output += ("\n" if output else "") + result.stderr
        
        output = output.strip()
        
        if not output:
            if result.returncode == 0:
                output = "(no output)"
            else:
                output = f"(exit code {result.returncode}, no output)"
        
        # Truncate very long output to prevent context explosion
        if len(output) > 3000:
            output = output[:2500] + f"\n... [truncated, {len(output)} total chars]"
        
        success = (result.returncode == 0)
        return success, output
        
    except subprocess.TimeoutExpired:
        return False, f"Command timed out after {timeout}s"
    except Exception as e:
        return False, f"Execution error: {str(e)}"


def _read_file(arguments: Dict) -> Tuple[bool, str]:
    """Read a file from the local filesystem."""
    path = arguments.get("path", "")
    if not path:
        return False, "No path specified"
    
    # Only allow local paths (no raspi: or gcp: prefix)
    if ":" in path and not path.startswith("/"):
        return False, "DELEGATE"  # Remote file, delegate to CLI
    
    # Expand user home
    path = os.path.expanduser(path)
    
    if not os.path.exists(path):
        return False, f"No such file or directory: {path}"
    
    if os.path.isdir(path):
        return False, f"Is a directory: {path}"
    
    max_lines = arguments.get("lines", None)
    
    try:
        with open(path, "r", errors="replace") as f:
            if max_lines:
                content = "".join(f.readline() for _ in range(max_lines))
            else:
                content = f.read()
        
        # Truncate
        if len(content) > 5000:
            content = content[:4500] + f"\n... [truncated, {len(content)} total chars]"
        
        return True, content.strip() if content.strip() else "(empty file)"
        
    except PermissionError:
        return False, f"Permission denied: {path}"
    except Exception as e:
        return False, f"Error reading {path}: {str(e)}"


def _search_files(arguments: Dict) -> Tuple[bool, str]:
    """Search for patterns in files (like grep -rn)."""
    pattern = arguments.get("pattern", "")
    if not pattern:
        return False, "No pattern specified"
    
    search_path = arguments.get("path", ".")
    
    # Only local paths
    if ":" in search_path and not search_path.startswith("/"):
        return False, "DELEGATE"
    
    search_path = os.path.expanduser(search_path)
    file_glob = arguments.get("glob", "*")
    
    try:
        # Use grep for efficiency
        cmd = f"grep -rn --include='{file_glob}' '{pattern}' {search_path} 2>/dev/null | head -50"
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=15)
        output = result.stdout.strip()
        
        if not output:
            return True, f"No matches found for '{pattern}' in {search_path}"
        
        if len(output) > 3000:
            output = output[:2500] + "\n... [truncated]"
        
        return True, output
        
    except Exception as e:
        return False, f"Search error: {str(e)}"


def _diagnose(arguments: Dict) -> Tuple[bool, str]:
    """Run basic diagnostics on the local system."""
    target = arguments.get("target", "local")
    
    if target != "local":
        return False, "DELEGATE"  # Remote diagnostics delegate to CLI
    
    checks = []
    
    # CPU load
    try:
        with open("/proc/loadavg", "r") as f:
            load = f.read().strip()
        checks.append(f"Load: {load}")
    except:
        pass
    
    # Memory
    try:
        result = subprocess.run("free -h | head -3", shell=True, capture_output=True, text=True, timeout=5)
        checks.append(f"Memory:\n{result.stdout.strip()}")
    except:
        pass
    
    # Disk
    try:
        result = subprocess.run("df -h / | tail -1", shell=True, capture_output=True, text=True, timeout=5)
        checks.append(f"Disk: {result.stdout.strip()}")
    except:
        pass
    
    # Docker
    try:
        result = subprocess.run("docker ps --format 'table {{.Names}}\t{{.Status}}' 2>/dev/null | head -10", shell=True, capture_output=True, text=True, timeout=5)
        if result.returncode == 0 and result.stdout.strip():
            checks.append(f"Docker:\n{result.stdout.strip()}")
    except:
        pass
    
    # Key services
    try:
        result = subprocess.run("systemctl list-units --type=service --state=running --no-pager | head -20", shell=True, capture_output=True, text=True, timeout=5)
        checks.append(f"Services:\n{result.stdout.strip()}")
    except:
        pass
    
    return True, "\n\n".join(checks) if checks else "Could not gather diagnostics"


def _robot_vision(arguments: Dict) -> Tuple[bool, str]:
    """Capture a snapshot from the robot's camera and analyze with moondream vision model."""
    import requests
    import base64
    import time

    ROBOT_PROXY = os.environ.get("ROBOT_PROXY_URL", "http://localhost:5002")
    OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:8200")
    VISION_MODEL = os.environ.get("VISION_MODEL", "moondream")

    question = arguments.get("question", "")
    if not question:
        question = "Describe what you see in this image in detail. What objects, people, furniture, or environment is visible? Be specific."

    # 1. Capture snapshot from robot camera
    t0 = time.time()
    try:
        resp = requests.get(f"{ROBOT_PROXY}/snapshot", timeout=10)
        resp.raise_for_status()
        img_data = resp.content
    except requests.ConnectionError:
        return False, "Robot camera unavailable -- cannot connect to robot proxy (localhost:5002). Is picar-forward.service running?"
    except Exception as e:
        return False, f"Error capturing snapshot from robot camera: {e}"

    t_capture = time.time() - t0

    if len(img_data) < 500:
        return False, f"Robot camera returned suspiciously small image ({len(img_data)} bytes). Camera may be offline."

    # 2. Encode to base64
    img_b64 = base64.b64encode(img_data).decode('utf-8')

    # 3. Analyze with vision model
    t1 = time.time()
    try:
        resp = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={
                "model": VISION_MODEL,
                "prompt": question,
                "images": [img_b64],
                "stream": False,
            },
            timeout=60,
        )
        resp.raise_for_status()
        result = resp.json()
    except requests.ConnectionError:
        return False, "Vision model unavailable -- cannot connect to Ollama (localhost:8200). Is tokioai-coder-tunnel.service running?"
    except Exception as e:
        return False, f"Error analyzing image with vision model: {e}"

    t_analysis = time.time() - t1

    description = result.get("response", "").strip()
    if not description:
        return False, "Vision model returned empty response."

    tokens = result.get("eval_count", 0)

    output = (
        f"[Robot Vision -- moondream]\n"
        f"Image: {len(img_data):,} bytes | Capture: {t_capture:.1f}s | Analysis: {t_analysis:.1f}s | Tokens: {tokens}\n"
        f"Question: {question}\n"
        f"Answer: {description}"
    )

    return True, output