"""
Device Scanner MCP Server
Detects WiFi networks, Ethernet connections, and USB devices.
Includes WiFi-hop: auto-switches WiFi to device, runs SSH commands, switches back.

All company-specific config is loaded from .env — see .env.example for setup.
"""

import json
import os
import re
import subprocess
import time
from pathlib import Path
from mcp.server.fastmcp import FastMCP

import paramiko

# ---------------------------------------------------------------------------
# Load .env — checks: current dir, home dir, then package dir
# ---------------------------------------------------------------------------
_SCRIPT_DIR = Path(__file__).resolve().parent
_ENV_SEARCH_PATHS = [
    Path.cwd() / ".env",                    # current working directory
    Path.home() / ".device-scanner-mcp.env", # user home
    Path.home() / ".env",                    # user home fallback
    _SCRIPT_DIR / ".env",                    # package directory
]
_ENV_FILE = next((p for p in _ENV_SEARCH_PATHS if p.exists()), _SCRIPT_DIR / ".env")


def _load_env(path: Path) -> None:
    """Minimal .env loader — no external dependency needed."""
    if not path.exists():
        return
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" not in line:
                continue
            key, _, value = line.partition("=")
            key = key.strip()
            value = value.strip()
            if not os.environ.get(key):
                os.environ[key] = value


_load_env(_ENV_FILE)


def _env(key: str, default: str = "") -> str:
    return os.environ.get(key, default)


def _env_list(key: str, default: str = "") -> list[str]:
    raw = _env(key, default)
    return [s.strip() for s in raw.split(",") if s.strip()]


def _env_int(key: str, default: int = 0) -> int:
    try:
        return int(_env(key, str(default)))
    except ValueError:
        return default


# ---------------------------------------------------------------------------
# Configuration from .env
# ---------------------------------------------------------------------------
COMPANY_NAME = _env("COMPANY_NAME", "MyCompany")
OFFICE_SSIDS = set(_env_list("OFFICE_SSIDS"))
DEVICE_PREFIXES = tuple(_env_list("DEVICE_PREFIXES"))
DEFAULT_SSH_USER = _env("DEFAULT_SSH_USER", "root")
DEFAULT_SSH_PASSWORD = _env("DEFAULT_SSH_PASSWORD", "")
DEFAULT_DEVICE_IP = _env("DEFAULT_DEVICE_IP", "10.0.0.10")
DEFAULT_SSH_PORT = _env_int("DEFAULT_SSH_PORT", 22)
DEFAULT_CMD_TIMEOUT = _env_int("DEFAULT_CMD_TIMEOUT", 300)
DEFAULT_SCRIPT_TIMEOUT = _env_int("DEFAULT_SCRIPT_TIMEOUT", 300)
WIFI_CONNECT_TIMEOUT = _env_int("WIFI_CONNECT_TIMEOUT", 15)
INTERNET_CHECK_TIMEOUT = _env_int("INTERNET_CHECK_TIMEOUT", 10)


_RESERVED_KEYS = {"DEVICE_PREFIXES", "DEVICE_ALIASES"}


def _load_device_profiles() -> dict:
    """Load DEVICE_* entries from env. Format: wifi_ip,eth_ip,usb_ip,user,password"""
    profiles = {}
    for key, value in os.environ.items():
        if key in _RESERVED_KEYS:
            continue
        if key.startswith("DEVICE_") and "," in value:
            name = key[len("DEVICE_"):]
            parts = value.split(",")
            if len(parts) >= 5:
                profiles[name] = {
                    k: v for k, v in {
                        "wifi_ip": parts[0].strip(),
                        "eth_ip": parts[1].strip(),
                        "usb_ip": parts[2].strip(),
                        "user": parts[3].strip(),
                        "password": parts[4].strip(),
                    }.items() if v
                }
    return profiles


def _load_device_aliases() -> dict:
    """Load device aliases. Format: alias1:prefix1,alias2:prefix2
    Returns {ALIAS_UPPER: SSID_PREFIX} mapping"""
    aliases = {}
    raw = _env("DEVICE_ALIASES", "")
    for pair in raw.split(","):
        pair = pair.strip()
        if ":" in pair:
            alias, prefix = pair.split(":", 1)
            aliases[alias.strip().upper()] = prefix.strip().upper()
    return aliases


DEVICE_PROFILES = _load_device_profiles()
DEVICE_ALIASES = _load_device_aliases()

# ---------------------------------------------------------------------------
# Build dynamic MCP instructions from config
# ---------------------------------------------------------------------------
_device_names = ", ".join(sorted(DEVICE_PROFILES.keys())) if DEVICE_PROFILES else "none configured"
_office_names = ", ".join(sorted(OFFICE_SSIDS)[:5]) if OFFICE_SSIDS else "none configured"
_prefix_list = ", ".join(DEVICE_PREFIXES[:6]) if DEVICE_PREFIXES else "none"
_alias_info = ", ".join(f"{a}={p}" for a, p in sorted(DEVICE_ALIASES.items())) if DEVICE_ALIASES else ""

mcp = FastMCP(
    "device-scanner",
    instructions=(
        f"Detects WiFi, Ethernet, and USB connections for {COMPANY_NAME} devices. "
        "Supports 3 connection types to devices: WiFi (hop), Ethernet (direct), USB (direct).\n\n"
        "CONNECTION TYPES:\n"
        "- WiFi: Use wifi_hop_* tools. These auto-switch WiFi to device, do work, switch back. "
        "REQUIRED because device WiFi has no internet.\n"
        "- Ethernet (IP 20.0.0.20): Use device_* tools with connection='ethernet'. "
        "Direct connection, no hop needed, internet stays connected.\n"
        "- USB (IP 30.0.0.30): Use device_* tools with connection='usb'. "
        "Direct connection, no hop needed, internet stays connected.\n\n"
        "TOOL SELECTION:\n"
        "- WiFi: wifi_hop_ssh, wifi_hop_ssh_batch, wifi_hop_scp_push, wifi_hop_scp_pull\n"
        "- Ethernet/USB: device_ssh, device_ssh_batch, device_scp_push, device_scp_pull\n\n"
        "CRITICAL WIFI RULES:\n"
        f"- NEVER use 'wifi_switch' for device SSIDs ({_prefix_list}...). Use wifi_hop_ssh.\n"
        f"- 'wifi_switch' is ONLY safe for office networks ({_office_names}).\n\n"
        f"DEVICES: {_device_names}.\n"
        f"ALIASES (product name = SSID prefix): {_alias_info}. "
        "When user says a product name, look for the aliased SSID prefix in scan results."
    ),
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _is_device_ssid(ssid: str) -> bool:
    return any(ssid.upper().startswith(p.upper()) for p in DEVICE_PREFIXES)


def _detect_device_type(ssid: str) -> str | None:
    """Detect device type from SSID, resolving aliases.
    Checks SSID against device profiles and resolves aliases."""
    ssid_upper = ssid.upper()
    # Direct match against profiles
    for prefix in DEVICE_PROFILES:
        if ssid_upper.startswith(prefix):
            return prefix
    # Check aliases: if SSID starts with an alias target prefix
    for alias, target_prefix in DEVICE_ALIASES.items():
        if ssid_upper.startswith(target_prefix):
            # Return the alias name if it has a profile, else the target
            if alias in DEVICE_PROFILES:
                return alias
            if target_prefix in DEVICE_PROFILES:
                return target_prefix
    return None


def _run_cmd(cmd: str, timeout: int = 15) -> str:
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, shell=True)
        return proc.stdout.strip()
    except subprocess.TimeoutExpired:
        return f"[timeout after {timeout}s]"
    except Exception as e:
        return f"[error: {e}]"


# ---------------------------------------------------------------------------
# WiFi scanning
# ---------------------------------------------------------------------------

def _scan_wifi() -> list[dict]:
    raw = _run_cmd("netsh wlan show networks mode=bssid", timeout=20)
    if raw.startswith("["):
        return [{"error": raw}]
    networks, current = [], {}
    for line in raw.splitlines():
        line = line.strip()
        if line.startswith("SSID") and ":" in line and "BSSID" not in line:
            if current:
                networks.append(current)
            current = {"ssid": line.split(":", 1)[1].strip()}
        elif line.startswith("Network type"):
            current["network_type"] = line.split(":", 1)[1].strip()
        elif line.startswith("Authentication"):
            current["authentication"] = line.split(":", 1)[1].strip()
        elif line.startswith("Encryption"):
            current["encryption"] = line.split(":", 1)[1].strip()
        elif line.startswith("BSSID"):
            current["bssid"] = line.split(":", 1)[1].strip()
        elif line.startswith("Signal"):
            current["signal"] = line.split(":", 1)[1].strip()
        elif line.startswith("Radio type"):
            current["radio_type"] = line.split(":", 1)[1].strip()
        elif line.startswith("Band"):
            current["band"] = line.split(":", 1)[1].strip()
        elif line.startswith("Channel") and "Utilization" not in line:
            current["channel"] = line.split(":", 1)[1].strip()
    if current:
        networks.append(current)
    return networks


@mcp.tool()
def scan_wifi() -> str:
    """List all visible WiFi networks nearby. Returns SSID, signal strength, band, channel, and authentication type for each network. Use this to find device WiFi SSIDs before connecting."""
    networks = _scan_wifi()
    if not networks:
        return json.dumps({"wifi_networks": [], "message": "No WiFi networks found"}, indent=2)
    return json.dumps({"wifi_networks": networks, "count": len(networks)}, indent=2)


# ---------------------------------------------------------------------------
# Network adapters
# ---------------------------------------------------------------------------

@mcp.tool()
def scan_network_adapters() -> str:
    """List all network adapters on this machine (Ethernet, WiFi, Bluetooth, virtual). Returns name, status (Up/Disconnected), MAC address, link speed, and IP addresses for each adapter."""
    raw = _run_cmd(
        'powershell -NoProfile -Command "'
        'Get-NetAdapter | Select-Object Name, InterfaceDescription, Status, '
        'MacAddress, LinkSpeed, MediaType | ConvertTo-Json"'
    )
    try:
        data = json.loads(raw)
        if isinstance(data, dict):
            data = [data]
        adapters = [{
            "name": a.get("Name", ""),
            "description": a.get("InterfaceDescription", ""),
            "status": a.get("Status", ""),
            "mac_address": a.get("MacAddress", ""),
            "link_speed": a.get("LinkSpeed", ""),
            "media_type": a.get("MediaType", ""),
        } for a in data]
    except json.JSONDecodeError:
        adapters = [{"error": raw[:200]}]

    raw_ip = _run_cmd(
        'powershell -NoProfile -Command "'
        "Get-NetIPAddress -AddressFamily IPv4 | Where-Object { \\$_.IPAddress -ne '127.0.0.1' } | "
        'Select-Object InterfaceAlias, IPAddress, PrefixLength | ConvertTo-Json"'
    )
    try:
        ip_data = json.loads(raw_ip)
        if isinstance(ip_data, dict):
            ip_data = [ip_data]
        ip_config = [{
            "interface": i.get("InterfaceAlias", ""),
            "ip_address": i.get("IPAddress", ""),
            "prefix_length": i.get("PrefixLength", ""),
        } for i in ip_data]
    except json.JSONDecodeError:
        ip_config = []

    return json.dumps({"network_adapters": adapters, "ip_configuration": ip_config, "count": len(adapters)}, indent=2)


# ---------------------------------------------------------------------------
# USB devices
# ---------------------------------------------------------------------------

@mcp.tool()
def scan_usb_devices() -> str:
    """List all USB devices currently connected to this machine. Returns device name, manufacturer, vendor/product IDs, and status."""
    raw = _run_cmd(
        'powershell -NoProfile -Command "'
        "Get-PnpDevice -Class USB | Where-Object { \\$_.Status -eq 'OK' } | "
        'Select-Object FriendlyName, InstanceId, Status, Manufacturer | ConvertTo-Json"'
    )
    try:
        data = json.loads(raw)
        if isinstance(data, dict):
            data = [data]
        devices = []
        for d in data:
            iid = d.get("InstanceId", "")
            vid = pid = ""
            m = re.search(r"VID_([0-9A-Fa-f]{4})", iid)
            if m: vid = m.group(1)
            m = re.search(r"PID_([0-9A-Fa-f]{4})", iid)
            if m: pid = m.group(1)
            devices.append({
                "name": d.get("FriendlyName", ""),
                "manufacturer": d.get("Manufacturer", ""),
                "status": d.get("Status", ""),
                "vendor_id": vid, "product_id": pid,
                "instance_id": iid,
            })
    except json.JSONDecodeError:
        devices = [{"error": raw[:200]}]

    return json.dumps({"usb_devices": devices, "count": len(devices)}, indent=2)


# ---------------------------------------------------------------------------
# Combined scan
# ---------------------------------------------------------------------------

@mcp.tool()
def scan_all_connections() -> str:
    """Run scan_wifi, scan_network_adapters, and scan_usb_devices together in one call. Use this for a complete overview of all connections."""
    return json.dumps({
        "wifi": json.loads(scan_wifi()),
        "network_adapters": json.loads(scan_network_adapters()),
        "usb": json.loads(scan_usb_devices()),
    }, indent=2)


# ---------------------------------------------------------------------------
# WiFi connect/status helpers
# ---------------------------------------------------------------------------

def _get_current_ssid() -> str:
    raw = _run_cmd("netsh wlan show interfaces", timeout=10)
    for line in raw.splitlines():
        line = line.strip()
        if line.startswith("SSID") and "BSSID" not in line:
            return line.split(":", 1)[1].strip()
    return ""


def _wifi_connect(ssid: str, timeout_sec: int = 0) -> tuple[bool, str]:
    timeout_sec = timeout_sec or WIFI_CONNECT_TIMEOUT
    result = _run_cmd(f'netsh wlan connect name="{ssid}"', timeout=10)
    if "completed successfully" not in result.lower():
        return False, f"Connect failed: {result}"
    for _ in range(timeout_sec * 2):
        time.sleep(0.5)
        if _get_current_ssid() == ssid:
            time.sleep(1)
            return True, f"Connected to {ssid}"
    return False, f"Timeout connecting to {ssid}"


def _wait_for_internet(timeout_sec: int = 0) -> bool:
    timeout_sec = timeout_sec or INTERNET_CHECK_TIMEOUT
    for _ in range(timeout_sec * 2):
        try:
            proc = subprocess.run(
                ["ping", "-n", "1", "-w", "1000", "8.8.8.8"],
                capture_output=True, timeout=3
            )
            if proc.returncode == 0:
                return True
        except Exception:
            pass
        time.sleep(0.5)
    return False


# ---------------------------------------------------------------------------
# SSH via paramiko
# ---------------------------------------------------------------------------

def _ssh_exec(host: str, user: str, password: str, commands: list[str],
              per_cmd_timeout: int = 0, port: int = 0) -> list[dict]:
    per_cmd_timeout = per_cmd_timeout or DEFAULT_CMD_TIMEOUT
    port = port or DEFAULT_SSH_PORT
    results = []
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(host, port=port, username=user, password=password,
                       timeout=8, allow_agent=False, look_for_keys=False)
        for cmd in commands:
            try:
                stdin, stdout, stderr = client.exec_command(cmd, timeout=per_cmd_timeout)
                exit_code = stdout.channel.recv_exit_status()
                results.append({
                    "command": cmd,
                    "stdout": stdout.read().decode("utf-8", errors="replace").strip(),
                    "stderr": stderr.read().decode("utf-8", errors="replace").strip(),
                    "exit_code": exit_code,
                })
            except Exception as e:
                results.append({"command": cmd, "stdout": "", "stderr": str(e), "exit_code": -1})
    except Exception as e:
        results.append({"command": commands[0] if commands else "", "stdout": "",
                        "stderr": f"SSH connection failed: {e}", "exit_code": -1})
    finally:
        client.close()
    return results


def _ssh_exec_script(host: str, user: str, password: str, script: str,
                     timeout: int = 0, port: int = 0) -> dict:
    timeout = timeout or DEFAULT_SCRIPT_TIMEOUT
    port = port or DEFAULT_SSH_PORT
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(host, port=port, username=user, password=password,
                       timeout=8, allow_agent=False, look_for_keys=False)
        stdin, stdout, stderr = client.exec_command("bash -s", timeout=timeout)
        stdin.write(script)
        stdin.channel.shutdown_write()
        exit_code = stdout.channel.recv_exit_status()
        return {
            "stdout": stdout.read().decode("utf-8", errors="replace").strip(),
            "stderr": stderr.read().decode("utf-8", errors="replace").strip(),
            "exit_code": exit_code,
        }
    except Exception as e:
        return {"stdout": "", "stderr": f"SSH failed: {e}", "exit_code": -1}
    finally:
        client.close()


# ---------------------------------------------------------------------------
# SFTP file transfer via paramiko
# ---------------------------------------------------------------------------

def _sftp_push(host: str, user: str, password: str, local_path: str,
               remote_path: str, port: int = 0) -> dict:
    """Push a file to device via SFTP."""
    port = port or DEFAULT_SSH_PORT
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(host, port=port, username=user, password=password,
                       timeout=8, allow_agent=False, look_for_keys=False)
        sftp = client.open_sftp()
        file_size = os.path.getsize(local_path)
        sftp.put(local_path, remote_path)
        # Verify
        remote_stat = sftp.stat(remote_path)
        sftp.close()
        return {
            "success": True,
            "local_path": local_path,
            "remote_path": remote_path,
            "bytes_transferred": file_size,
            "remote_size": remote_stat.st_size,
        }
    except Exception as e:
        return {"success": False, "error": str(e)}
    finally:
        client.close()


def _sftp_pull(host: str, user: str, password: str, remote_path: str,
               local_path: str, port: int = 0) -> dict:
    """Pull a file from device via SFTP."""
    port = port or DEFAULT_SSH_PORT
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(host, port=port, username=user, password=password,
                       timeout=8, allow_agent=False, look_for_keys=False)
        sftp = client.open_sftp()
        sftp.get(remote_path, local_path)
        sftp.close()
        local_size = os.path.getsize(local_path)
        return {
            "success": True,
            "remote_path": remote_path,
            "local_path": local_path,
            "bytes_transferred": local_size,
        }
    except Exception as e:
        return {"success": False, "error": str(e)}
    finally:
        client.close()


# ---------------------------------------------------------------------------
# Connection type resolver
# ---------------------------------------------------------------------------

def _resolve_device(device_type: str, connection: str = "wifi") -> dict:
    """Resolve device credentials from type + connection method.
    connection: 'wifi', 'ethernet'/'eth', or 'usb'
    Returns {ip, user, password} or {error}."""
    device_type = device_type.upper()
    # Resolve alias (e.g. product name -> SSID prefix)
    if device_type in DEVICE_ALIASES:
        resolved = DEVICE_ALIASES[device_type]
        # Use resolved name's profile if it exists, else try original
        if resolved in DEVICE_PROFILES:
            device_type = resolved
    if device_type not in DEVICE_PROFILES:
        return {"error": f"Unknown device type: {device_type}. Known: {list(DEVICE_PROFILES.keys())}"}
    p = DEVICE_PROFILES[device_type]
    conn = connection.lower().strip()
    ip_key = {"wifi": "wifi_ip", "ethernet": "eth_ip", "eth": "eth_ip", "usb": "usb_ip"}.get(conn)
    if not ip_key:
        return {"error": f"Unknown connection type: {connection}. Use wifi, ethernet, or usb."}
    ip = p.get(ip_key, "")
    if not ip:
        return {"error": f"{device_type} has no {connection} IP configured."}
    return {"ip": ip, "user": p.get("user", DEFAULT_SSH_USER), "password": p.get("password", DEFAULT_SSH_PASSWORD)}


# ---------------------------------------------------------------------------
# Direct SSH/SCP tools (for Ethernet & USB — no WiFi hop needed)
# ---------------------------------------------------------------------------

@mcp.tool()
def device_ssh(device_type: str, commands: list[str],
               connection: str = "ethernet", device_ip: str = "",
               ssh_user: str = "", ssh_password: str = "",
               per_cmd_timeout: int = 0) -> str:
    """
    Run one or more shell commands on a device over Ethernet or USB.
    Each command runs sequentially in a single SSH session. Use this for quick checks like 'uname -a', 'ls /tmp', 'df -h'.
    For multi-line scripts with pipes or variables, use device_ssh_batch instead.
    Does NOT require WiFi hop — the device is directly reachable via wired connection.

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        commands: List of shell commands — e.g. ["uname -a", "ls /tmp", "free -h"]
        connection: How the device is connected — "ethernet" (IP 20.0.0.20) or "usb" (IP 30.0.0.30)
        device_ip: Override the auto-detected IP address
        ssh_user: Override the auto-detected SSH username
        ssh_password: Override the auto-detected SSH password
        per_cmd_timeout: Max seconds to wait per command (default 300)
    """
    per_cmd_timeout = per_cmd_timeout or DEFAULT_CMD_TIMEOUT

    resolved = _resolve_device(device_type, connection)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)

    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    t0 = time.time()
    results = _ssh_exec(ip, user, pw, commands, per_cmd_timeout)

    return json.dumps({
        "device_type": device_type.upper(),
        "connection": connection,
        "device_ip": ip,
        "command_results": results,
        "duration": f"{time.time() - t0:.1f}s",
    }, indent=2)


@mcp.tool()
def device_ssh_batch(device_type: str, script: str,
                     connection: str = "ethernet", device_ip: str = "",
                     ssh_user: str = "", ssh_password: str = "",
                     timeout: int = 0) -> str:
    """
    Run a multi-line bash script on a device over Ethernet or USB.
    The entire script is sent as stdin to 'bash -s' in a single SSH session.
    Use this when commands depend on each other (variables, pipes, loops, conditionals).
    For simple independent commands, use device_ssh instead.

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        script: Multi-line bash script — e.g. "cd /tmp\\nexport LD_LIBRARY_PATH=.\\n./my_test"
        connection: "ethernet" or "usb"
        device_ip: Override the auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
        timeout: Max seconds for the entire script (default 300)
    """
    timeout = timeout or DEFAULT_SCRIPT_TIMEOUT

    resolved = _resolve_device(device_type, connection)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)

    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    t0 = time.time()
    result = _ssh_exec_script(ip, user, pw, script, timeout)

    return json.dumps({
        "device_type": device_type.upper(),
        "connection": connection,
        "device_ip": ip,
        "script_result": result,
        "duration": f"{time.time() - t0:.1f}s",
    }, indent=2)


@mcp.tool()
def device_scp_push(device_type: str, local_path: str, remote_path: str,
                    connection: str = "ethernet", device_ip: str = "",
                    ssh_user: str = "", ssh_password: str = "") -> str:
    """
    Upload a single file from this machine to a device over Ethernet or USB.
    Uses SFTP. The remote directory must already exist — use device_ssh to create it first if needed.
    For uploading over WiFi, use wifi_hop_scp_push instead.

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        local_path: Full path to the local file — e.g. "C:/temp/firmware.bin"
        remote_path: Full destination path on device — e.g. "/tmp/firmware.bin"
        connection: "ethernet" or "usb"
        device_ip: Override the auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
    """
    if not os.path.exists(local_path):
        return json.dumps({"error": f"File not found: {local_path}"}, indent=2)

    resolved = _resolve_device(device_type, connection)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)

    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    t0 = time.time()
    result = _sftp_push(ip, user, pw, local_path, remote_path)
    result["duration"] = f"{time.time() - t0:.1f}s"
    result["connection"] = connection

    return json.dumps(result, indent=2)


@mcp.tool()
def device_scp_pull(device_type: str, remote_path: str, local_path: str,
                    connection: str = "ethernet", device_ip: str = "",
                    ssh_user: str = "", ssh_password: str = "") -> str:
    """
    Download a single file from a device to this machine over Ethernet or USB.
    Uses SFTP. For downloading over WiFi, use wifi_hop_scp_pull instead.

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        remote_path: Full path of the file on device — e.g. "/tmp/logs/app.log"
        local_path: Full local destination path — e.g. "C:/temp/app.log"
        connection: "ethernet" or "usb"
        device_ip: Override the auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
    """
    resolved = _resolve_device(device_type, connection)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)

    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    t0 = time.time()
    result = _sftp_pull(ip, user, pw, remote_path, local_path)
    result["duration"] = f"{time.time() - t0:.1f}s"
    result["connection"] = connection

    return json.dumps(result, indent=2)


# ---------------------------------------------------------------------------
# WiFi Hop tools (WiFi only — needs network switch)
# ---------------------------------------------------------------------------

@mcp.tool()
def wifi_hop_scp_push(device_ssid: str, local_path: str, remote_path: str,
                      home_ssid: str = "", device_ip: str = "",
                      ssh_user: str = "", ssh_password: str = "") -> str:
    """
    Upload a single file to a device over WiFi. Automatically switches WiFi to the device network, uploads via SFTP, then switches back to office WiFi.
    The remote directory must already exist. For Ethernet/USB use device_scp_push instead (no hop needed).

    Args:
        device_ssid: WiFi SSID of the device — e.g. "DEVICE_ABC123"
        local_path: Full local file path — e.g. "C:/temp/firmware.bin"
        remote_path: Full remote destination — e.g. "/tmp/firmware.bin"
        home_ssid: Office WiFi to return to (auto-detected if empty)
        device_ip: Override auto-detected IP (default 10.0.0.10)
        ssh_user: SSH user (auto-detected)
        ssh_password: SSH password (auto-detected)
    """
    log = []
    original_ssid = home_ssid or _get_current_ssid()
    if not original_ssid:
        return json.dumps({"error": "Cannot detect current WiFi."}, indent=2)

    if not os.path.exists(local_path):
        return json.dumps({"error": f"File not found: {local_path}"}, indent=2)

    device_type = _detect_device_type(device_ssid)
    if device_type and device_type in DEVICE_PROFILES:
        p = DEVICE_PROFILES[device_type]
        device_ip = device_ip or p.get("wifi_ip", DEFAULT_DEVICE_IP)
        ssh_user = ssh_user or p.get("user", DEFAULT_SSH_USER)
        ssh_password = ssh_password or p.get("password", DEFAULT_SSH_PASSWORD)
    else:
        device_ip = device_ip or DEFAULT_DEVICE_IP
        ssh_user = ssh_user or DEFAULT_SSH_USER
        ssh_password = ssh_password or DEFAULT_SSH_PASSWORD

    t0 = time.time()
    ok, msg = _wifi_connect(device_ssid)
    if not ok:
        _wifi_connect(original_ssid)
        return json.dumps({"error": msg, "home_restored": True}, indent=2)
    log.append(f"Device WiFi: {time.time() - t0:.1f}s")

    t1 = time.time()
    result = _sftp_push(device_ip, ssh_user, ssh_password, local_path, remote_path)
    log.append(f"SFTP push: {time.time() - t1:.1f}s")

    t2 = time.time()
    ok, _ = _wifi_connect(original_ssid)
    if not ok:
        subprocess.run("netsh wlan disconnect", shell=True, capture_output=True, timeout=5)
        time.sleep(1)
        ok, _ = _wifi_connect(original_ssid)
    log.append(f"Home WiFi: {time.time() - t2:.1f}s")

    if ok:
        t3 = time.time()
        inet = _wait_for_internet()
        log.append(f"Internet: {'OK' if inet else 'FAILED'} ({time.time() - t3:.1f}s)")

    log.append(f"Total: {time.time() - t0:.1f}s")
    return json.dumps({"transfer": result, "log": log, "home_restored": ok}, indent=2)


@mcp.tool()
def wifi_hop_scp_pull(device_ssid: str, remote_path: str, local_path: str,
                      home_ssid: str = "", device_ip: str = "",
                      ssh_user: str = "", ssh_password: str = "") -> str:
    """
    Download a single file from a device over WiFi. Automatically switches WiFi to the device network, downloads via SFTP, then switches back to office WiFi.
    For Ethernet/USB use device_scp_pull instead (no hop needed).

    Args:
        device_ssid: WiFi SSID of the device — e.g. "DEVICE_XYZ789"
        remote_path: Full path of file on device — e.g. "/tmp/logs/crash.log"
        local_path: Full local destination — e.g. "C:/temp/crash.log"
        home_ssid: Office WiFi to return to (auto-detected if empty)
        device_ip: Override auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
    """
    log = []
    original_ssid = home_ssid or _get_current_ssid()
    if not original_ssid:
        return json.dumps({"error": "Cannot detect current WiFi."}, indent=2)

    device_type = _detect_device_type(device_ssid)
    if device_type and device_type in DEVICE_PROFILES:
        p = DEVICE_PROFILES[device_type]
        device_ip = device_ip or p.get("wifi_ip", DEFAULT_DEVICE_IP)
        ssh_user = ssh_user or p.get("user", DEFAULT_SSH_USER)
        ssh_password = ssh_password or p.get("password", DEFAULT_SSH_PASSWORD)
    else:
        device_ip = device_ip or DEFAULT_DEVICE_IP
        ssh_user = ssh_user or DEFAULT_SSH_USER
        ssh_password = ssh_password or DEFAULT_SSH_PASSWORD

    t0 = time.time()
    ok, msg = _wifi_connect(device_ssid)
    if not ok:
        _wifi_connect(original_ssid)
        return json.dumps({"error": msg, "home_restored": True}, indent=2)
    log.append(f"Device WiFi: {time.time() - t0:.1f}s")

    t1 = time.time()
    result = _sftp_pull(device_ip, ssh_user, ssh_password, remote_path, local_path)
    log.append(f"SFTP pull: {time.time() - t1:.1f}s")

    t2 = time.time()
    ok, _ = _wifi_connect(original_ssid)
    if not ok:
        subprocess.run("netsh wlan disconnect", shell=True, capture_output=True, timeout=5)
        time.sleep(1)
        ok, _ = _wifi_connect(original_ssid)
    log.append(f"Home WiFi: {time.time() - t2:.1f}s")

    if ok:
        t3 = time.time()
        inet = _wait_for_internet()
        log.append(f"Internet: {'OK' if inet else 'FAILED'} ({time.time() - t3:.1f}s)")

    log.append(f"Total: {time.time() - t0:.1f}s")
    return json.dumps({"transfer": result, "log": log, "home_restored": ok}, indent=2)


@mcp.tool()
def wifi_hop_ssh(device_ssid: str, commands: list[str],
                 home_ssid: str = "", device_ip: str = "",
                 ssh_user: str = "", ssh_password: str = "",
                 per_cmd_timeout: int = 0) -> str:
    """
    Run shell commands on a device over WiFi. Automatically switches WiFi to the device, runs commands via SSH, then switches back to office WiFi.
    For Ethernet/USB use device_ssh instead (no hop needed, faster).

    Args:
        device_ssid: WiFi SSID of the device — e.g. "DEVICE_ABC123", "DEVICE_XYZ789"
        commands: List of shell commands — e.g. ["uname -a", "ls /tmp", "df -h"]
        home_ssid: Office WiFi to return to (auto-detected if empty)
        device_ip: Override auto-detected IP (default 10.0.0.10)
        ssh_user: Override SSH username
        ssh_password: Override SSH password
        per_cmd_timeout: Max seconds per command (default 300)
    """
    per_cmd_timeout = per_cmd_timeout or DEFAULT_CMD_TIMEOUT
    log = []

    original_ssid = home_ssid or _get_current_ssid()
    if not original_ssid:
        return json.dumps({"error": "Cannot detect current WiFi."}, indent=2)
    log.append(f"Home: {original_ssid}")

    # Auto-detect credentials from device type
    device_type = _detect_device_type(device_ssid)
    if device_type and device_type in DEVICE_PROFILES:
        p = DEVICE_PROFILES[device_type]
        device_ip = device_ip or p.get("wifi_ip", DEFAULT_DEVICE_IP)
        ssh_user = ssh_user or p.get("user", DEFAULT_SSH_USER)
        ssh_password = ssh_password or p.get("password", DEFAULT_SSH_PASSWORD)
    else:
        device_ip = device_ip or DEFAULT_DEVICE_IP
        ssh_user = ssh_user or DEFAULT_SSH_USER
        ssh_password = ssh_password or DEFAULT_SSH_PASSWORD
        if not ssh_password:
            return json.dumps({"error": f"Unknown device '{device_ssid}'. Provide ssh_password."}, indent=2)

    # Switch to device
    t0 = time.time()
    ok, msg = _wifi_connect(device_ssid)
    if not ok:
        _wifi_connect(original_ssid)
        return json.dumps({"error": msg, "log": log, "home_restored": True}, indent=2)
    log.append(f"Device WiFi: {time.time() - t0:.1f}s")

    # Run commands
    t1 = time.time()
    cmd_results = _ssh_exec(device_ip, ssh_user, ssh_password, commands, per_cmd_timeout)
    log.append(f"SSH commands: {time.time() - t1:.1f}s")

    # Switch back
    t2 = time.time()
    ok, msg = _wifi_connect(original_ssid)
    home_restored = ok
    if not ok:
        subprocess.run("netsh wlan disconnect", shell=True, capture_output=True, timeout=5)
        time.sleep(1)
        ok, _ = _wifi_connect(original_ssid)
        home_restored = ok
    log.append(f"Home WiFi: {time.time() - t2:.1f}s")

    if home_restored:
        t3 = time.time()
        inet = _wait_for_internet()
        log.append(f"Internet: {'OK' if inet else 'FAILED'} ({time.time() - t3:.1f}s)")

    log.append(f"Total: {time.time() - t0:.1f}s")

    return json.dumps({
        "device_ssid": device_ssid,
        "device_type": device_type,
        "device_ip": device_ip,
        "home_ssid": original_ssid,
        "home_restored": home_restored,
        "command_results": cmd_results,
        "log": log,
    }, indent=2)


@mcp.tool()
def wifi_hop_ssh_batch(device_ssid: str, script: str,
                       home_ssid: str = "", device_ip: str = "",
                       ssh_user: str = "", ssh_password: str = "",
                       timeout: int = 0) -> str:
    """
    Run a multi-line bash script on a device over WiFi. Automatically switches WiFi to the device, runs the script via SSH, then switches back.
    For Ethernet/USB use device_ssh_batch instead (no hop needed, faster).

    Args:
        device_ssid: WiFi SSID of the device — e.g. "DEVICE_ABC123"
        script: Multi-line bash script — e.g. "cd /tmp\\nexport LD_LIBRARY_PATH=.\\n./run_tests.sh"
        home_ssid: Office WiFi to return to (auto-detected if empty)
        device_ip: Override auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
        timeout: Max seconds for the entire script (default 300)
    """
    timeout = timeout or DEFAULT_SCRIPT_TIMEOUT
    log = []

    original_ssid = home_ssid or _get_current_ssid()
    if not original_ssid:
        return json.dumps({"error": "Cannot detect current WiFi."}, indent=2)

    device_type = _detect_device_type(device_ssid)
    if device_type and device_type in DEVICE_PROFILES:
        p = DEVICE_PROFILES[device_type]
        device_ip = device_ip or p.get("wifi_ip", DEFAULT_DEVICE_IP)
        ssh_user = ssh_user or p.get("user", DEFAULT_SSH_USER)
        ssh_password = ssh_password or p.get("password", DEFAULT_SSH_PASSWORD)
    else:
        device_ip = device_ip or DEFAULT_DEVICE_IP
        ssh_user = ssh_user or DEFAULT_SSH_USER
        ssh_password = ssh_password or DEFAULT_SSH_PASSWORD
        if not ssh_password:
            return json.dumps({"error": "Unknown device. Provide ssh_password."}, indent=2)

    t0 = time.time()
    ok, msg = _wifi_connect(device_ssid)
    if not ok:
        _wifi_connect(original_ssid)
        return json.dumps({"error": msg, "home_restored": True}, indent=2)
    log.append(f"Device WiFi: {time.time() - t0:.1f}s")

    t1 = time.time()
    result = _ssh_exec_script(device_ip, ssh_user, ssh_password, script, timeout)
    log.append(f"Script: {time.time() - t1:.1f}s")

    t2 = time.time()
    ok, _ = _wifi_connect(original_ssid)
    if not ok:
        subprocess.run("netsh wlan disconnect", shell=True, capture_output=True, timeout=5)
        time.sleep(1)
        ok, _ = _wifi_connect(original_ssid)
    log.append(f"Home WiFi: {time.time() - t2:.1f}s")

    if ok:
        t3 = time.time()
        inet = _wait_for_internet()
        log.append(f"Internet: {'OK' if inet else 'FAILED'} ({time.time() - t3:.1f}s)")

    log.append(f"Total: {time.time() - t0:.1f}s")

    return json.dumps({
        "device_ssid": device_ssid,
        "device_type": device_type,
        "home_ssid": original_ssid,
        "home_restored": ok,
        "script_result": result,
        "log": log,
    }, indent=2)


# ---------------------------------------------------------------------------
# Directory upload — tar a local directory and upload the .tar.gz to device
# ---------------------------------------------------------------------------

@mcp.tool()
def device_upload_directory(device_type: str, local_dir: str,
                           remote_path: str = "/tmp/upload.tar.gz",
                           connection: str = "ethernet", device_ssid: str = "",
                           file_glob: str = "", device_ip: str = "",
                           ssh_user: str = "", ssh_password: str = "") -> str:
    """
    Compress a local directory into a .tar.gz and upload it to a device.
    Does ONE job: tar locally + upload. Does NOT extract — use device_extract_archive to extract.
    Works over Ethernet, USB, or WiFi (auto-hops for WiFi).

    Typical workflow to deploy files:
      1. device_upload_directory → uploads .tar.gz to device
      2. device_extract_archive → extracts .tar.gz on device

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        local_dir: Local directory to compress — e.g. "C:/temp/binaries"
        remote_path: Destination .tar.gz path on device — e.g. "/tmp/binaries.tar.gz"
        connection: "ethernet", "usb", or "wifi"
        device_ssid: Required only for WiFi — the device SSID to hop to
        file_glob: Optional filter — e.g. "*.so" to only include .so files, or "test_*" for test binaries
        device_ip: Override the auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
    """
    import glob as glob_mod
    import tarfile
    import tempfile

    if not os.path.isdir(local_dir):
        return json.dumps({"error": f"Directory not found: {local_dir}"}, indent=2)

    conn = connection.lower().strip()
    resolved = _resolve_device(device_type, conn)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)
    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    log = []
    t0 = time.time()

    # Step 1: Create tarball locally
    tmp_tar = tempfile.mktemp(suffix=".tar.gz")
    try:
        with tarfile.open(tmp_tar, "w:gz") as tar:
            if file_glob:
                files = glob_mod.glob(os.path.join(local_dir, file_glob))
                for f in files:
                    tar.add(f, arcname=os.path.basename(f))
            else:
                for item in os.listdir(local_dir):
                    full = os.path.join(local_dir, item)
                    if os.path.isfile(full):
                        tar.add(full, arcname=item)
        tar_size = os.path.getsize(tmp_tar)
        log.append(f"Tar created: {tar_size / 1024 / 1024:.1f}MB")
    except Exception as e:
        return json.dumps({"error": f"Tar failed: {e}"}, indent=2)

    # Step 2: WiFi hop if needed
    original_ssid = None
    if conn == "wifi":
        if not device_ssid:
            os.unlink(tmp_tar)
            return json.dumps({"error": "device_ssid required for WiFi."}, indent=2)
        original_ssid = _get_current_ssid()
        ok, msg = _wifi_connect(device_ssid)
        if not ok:
            _wifi_connect(original_ssid)
            os.unlink(tmp_tar)
            return json.dumps({"error": msg, "home_restored": True}, indent=2)

    # Step 3: Upload via SFTP
    t1 = time.time()
    push_result = _sftp_push(ip, user, pw, tmp_tar, remote_path)
    os.unlink(tmp_tar)
    log.append(f"Upload: {time.time() - t1:.1f}s")

    # Hop back if WiFi
    home_restored = True
    if original_ssid:
        ok, _ = _wifi_connect(original_ssid)
        home_restored = ok
        if ok:
            _wait_for_internet()

    log.append(f"Total: {time.time() - t0:.1f}s")

    return json.dumps({
        "device_type": device_type.upper(),
        "connection": connection,
        "local_dir": local_dir,
        "remote_path": remote_path,
        "bytes_transferred": tar_size,
        "upload_success": push_result.get("success", False),
        "home_restored": home_restored,
        "log": log,
    }, indent=2)


# ---------------------------------------------------------------------------
# Extract archive on device
# ---------------------------------------------------------------------------

@mcp.tool()
def device_extract_archive(device_type: str, archive_path: str,
                           extract_dir: str = "/tmp",
                           connection: str = "ethernet", device_ssid: str = "",
                           device_ip: str = "", ssh_user: str = "",
                           ssh_password: str = "",
                           delete_archive: bool = True) -> str:
    """
    Extract a .tar.gz archive that is already on a device.
    Does ONE job: extract. Use device_upload_directory or device_scp_push to upload the archive first.
    Works over Ethernet, USB, or WiFi (auto-hops for WiFi).

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        archive_path: Path to .tar.gz on device — e.g. "/tmp/binaries.tar.gz"
        extract_dir: Directory to extract into — e.g. "/tmp/binaries"
        connection: "ethernet", "usb", or "wifi"
        device_ssid: Required only for WiFi
        device_ip: Override the auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
        delete_archive: Delete the .tar.gz after extraction (default true)
    """
    conn = connection.lower().strip()
    resolved = _resolve_device(device_type, conn)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)
    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    rm_cmd = f" && rm {archive_path}" if delete_archive else ""
    cmd = f"mkdir -p {extract_dir} && cd {extract_dir} && tar xzf {archive_path}{rm_cmd} && ls -la"

    original_ssid = None
    if conn == "wifi":
        if not device_ssid:
            return json.dumps({"error": "device_ssid required for WiFi."}, indent=2)
        original_ssid = _get_current_ssid()
        ok, msg = _wifi_connect(device_ssid)
        if not ok:
            _wifi_connect(original_ssid)
            return json.dumps({"error": msg}, indent=2)

    results = _ssh_exec(ip, user, pw, [cmd], per_cmd_timeout=60)

    home_restored = True
    if original_ssid:
        ok, _ = _wifi_connect(original_ssid)
        home_restored = ok
        if ok:
            _wait_for_internet()

    return json.dumps({
        "device_type": device_type.upper(),
        "archive_path": archive_path,
        "extract_dir": extract_dir,
        "deleted_archive": delete_archive,
        "result": results[0] if results else {},
        "home_restored": home_restored,
    }, indent=2)


# ---------------------------------------------------------------------------
# HIGH VALUE: device_health — quick device status
# ---------------------------------------------------------------------------

@mcp.tool()
def device_health(device_type: str, connection: str = "ethernet",
                  device_ssid: str = "", device_ip: str = "",
                  ssh_user: str = "", ssh_password: str = "") -> str:
    """
    Get a quick health snapshot of a device: CPU info, memory usage, disk space, uptime, temperature, load average, top processes, and network interfaces.
    Works over Ethernet, USB, or WiFi (auto-hops for WiFi).

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        connection: "ethernet", "usb", or "wifi"
        device_ssid: Required only for WiFi — the device SSID
        device_ip: Override auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
    """
    conn = connection.lower().strip()
    resolved = _resolve_device(device_type, conn)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)
    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    health_script = """
echo "===UNAME==="
uname -a
echo "===UPTIME==="
uptime
echo "===CPU==="
cat /proc/cpuinfo | grep -c processor
cat /proc/cpuinfo | grep "model name" | head -1 2>/dev/null || echo "N/A"
echo "===MEMORY==="
free -h
echo "===DISK==="
df -h / /tmp 2>/dev/null
echo "===TEMPERATURE==="
cat /sys/class/thermal/thermal_zone*/temp 2>/dev/null | while read t; do echo "$((t/1000))C"; done
echo "===LOAD==="
cat /proc/loadavg
echo "===TOP_PROCESSES==="
ps aux --sort=-%cpu 2>/dev/null | head -8 || ps -ef | head -8
echo "===NETWORK==="
ip addr show 2>/dev/null | grep -E "inet |state" | head -10
"""

    original_ssid = None
    if conn == "wifi":
        if not device_ssid:
            return json.dumps({"error": "device_ssid required for WiFi."}, indent=2)
        original_ssid = _get_current_ssid()
        ok, msg = _wifi_connect(device_ssid)
        if not ok:
            _wifi_connect(original_ssid)
            return json.dumps({"error": msg}, indent=2)

    t0 = time.time()
    result = _ssh_exec_script(ip, user, pw, health_script, timeout=30)

    if original_ssid:
        _wifi_connect(original_ssid)
        _wait_for_internet()

    # Parse sections
    raw = result.get("stdout", "")
    sections = {}
    current_key = None
    current_lines = []
    for line in raw.splitlines():
        if line.startswith("===") and line.endswith("==="):
            if current_key:
                sections[current_key] = "\n".join(current_lines).strip()
            current_key = line.strip("=")
            current_lines = []
        else:
            current_lines.append(line)
    if current_key:
        sections[current_key] = "\n".join(current_lines).strip()

    return json.dumps({
        "device_type": device_type.upper(),
        "connection": connection,
        "device_ip": ip,
        "health": sections,
        "duration": f"{time.time() - t0:.1f}s",
    }, indent=2)


# ---------------------------------------------------------------------------
# HIGH VALUE: device_ping — check reachability
# ---------------------------------------------------------------------------

@mcp.tool()
def device_ping(device_type: str, connection: str = "ethernet",
                device_ip: str = "") -> str:
    """
    Ping a device to check if it is reachable on the network. Use this before SSH/SCP to avoid long timeouts on offline devices.
    Only works for Ethernet/USB (direct connections). Cannot ping device WiFi without hopping.

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        connection: "ethernet" or "usb"
        device_ip: Override auto-detected IP
    """
    resolved = _resolve_device(device_type, connection)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)
    ip = device_ip or resolved["ip"]

    try:
        proc = subprocess.run(
            ["ping", "-n", "3", "-w", "1000", ip],
            capture_output=True, text=True, timeout=10
        )
        reachable = proc.returncode == 0
        return json.dumps({
            "device_type": device_type.upper(),
            "connection": connection,
            "ip": ip,
            "reachable": reachable,
            "output": proc.stdout.strip()[-300:],
        }, indent=2)
    except Exception as e:
        return json.dumps({"ip": ip, "reachable": False, "error": str(e)}, indent=2)


# ---------------------------------------------------------------------------
# HIGH VALUE: device_reboot — reboot and wait for recovery
# ---------------------------------------------------------------------------

@mcp.tool()
def device_reboot(device_type: str, connection: str = "ethernet",
                  device_ssid: str = "", device_ip: str = "",
                  ssh_user: str = "", ssh_password: str = "",
                  wait_timeout: int = 120) -> str:
    """
    Reboot a device and optionally wait until it comes back online.
    Sends 'reboot' via SSH, then polls with ping+SSH until the device responds again.
    Works over Ethernet, USB, or WiFi (auto-hops for WiFi).

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        connection: "ethernet", "usb", or "wifi"
        device_ssid: Required only for WiFi
        device_ip: Override auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
        wait_timeout: Max seconds to wait for recovery (0 = fire and forget, don't wait)
    """
    conn = connection.lower().strip()
    resolved = _resolve_device(device_type, conn)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)
    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    log = []

    # WiFi hop if needed
    original_ssid = None
    if conn == "wifi":
        if not device_ssid:
            return json.dumps({"error": "device_ssid required for WiFi."}, indent=2)
        original_ssid = _get_current_ssid()
        ok, msg = _wifi_connect(device_ssid)
        if not ok:
            _wifi_connect(original_ssid)
            return json.dumps({"error": msg}, indent=2)

    # Send reboot command
    t0 = time.time()
    reboot_result = _ssh_exec(ip, user, pw, ["nohup sh -c 'sleep 1 && reboot' &>/dev/null &"], per_cmd_timeout=10)
    log.append("Reboot command sent")

    # Hop back immediately if WiFi
    if original_ssid:
        time.sleep(2)
        _wifi_connect(original_ssid)
        _wait_for_internet()
        log.append("WiFi restored")

    if wait_timeout <= 0:
        return json.dumps({
            "device_type": device_type.upper(),
            "rebooted": True,
            "waited_for_recovery": False,
            "log": log,
        }, indent=2)

    # Wait for device to go down then come back
    log.append(f"Waiting up to {wait_timeout}s for recovery...")
    time.sleep(10)  # Give it time to actually go down

    # For WiFi, we need to hop again to check
    if conn == "wifi":
        recovered = False
        for attempt in range(0, wait_timeout - 10, 15):
            time.sleep(15)
            original_ssid = _get_current_ssid()
            ok, _ = _wifi_connect(device_ssid)
            if ok:
                try:
                    check = _ssh_exec(ip, user, pw, ["uptime"], per_cmd_timeout=10)
                    if check and check[0].get("exit_code") == 0:
                        recovered = True
                        log.append(f"Device back online ({time.time() - t0:.0f}s)")
                        _wifi_connect(original_ssid)
                        _wait_for_internet()
                        break
                except Exception:
                    pass
                _wifi_connect(original_ssid)
                _wait_for_internet()
    else:
        # Ethernet/USB — just ping until it's back
        recovered = False
        deadline = time.time() + wait_timeout - 10
        while time.time() < deadline:
            try:
                proc = subprocess.run(
                    ["ping", "-n", "1", "-w", "1000", ip],
                    capture_output=True, timeout=5
                )
                if proc.returncode == 0:
                    time.sleep(5)  # Give SSH time to start
                    check = _ssh_exec(ip, user, pw, ["uptime"], per_cmd_timeout=10)
                    if check and check[0].get("exit_code") == 0:
                        recovered = True
                        log.append(f"Device back online ({time.time() - t0:.0f}s)")
                        break
            except Exception:
                pass
            time.sleep(5)

    if not recovered:
        log.append(f"Device did NOT recover within {wait_timeout}s")

    return json.dumps({
        "device_type": device_type.upper(),
        "connection": connection,
        "rebooted": True,
        "recovered": recovered,
        "total_time": f"{time.time() - t0:.0f}s",
        "log": log,
    }, indent=2)


# ---------------------------------------------------------------------------
# HIGH VALUE: device_run_tests — run GTest binaries, parse results
# ---------------------------------------------------------------------------

@mcp.tool()
def device_run_tests(device_type: str, test_binaries: list[str],
                     test_dir: str = "/tmp", connection: str = "ethernet",
                     device_ssid: str = "", device_ip: str = "",
                     ssh_user: str = "", ssh_password: str = "",
                     ld_library_path: str = "", timeout: int = 0) -> str:
    """
    Run GTest/executable test binaries on a device and return structured pass/fail results.
    Automatically parses Google Test output to extract passed/failed counts per binary.
    Returns a summary table with per-binary results and overall totals.
    Works over Ethernet, USB, or WiFi (auto-hops for WiFi).

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        test_binaries: List of binary names to run — e.g. ["smoke_test", "test_iox_wrapper", "iox_flatbuf_tests"]
        test_dir: Directory on device where the binaries are located — e.g. "/tmp" or "/tmp/middleware4"
        connection: "ethernet", "usb", or "wifi"
        device_ssid: Required only for WiFi
        device_ip: Override auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
        ld_library_path: Set LD_LIBRARY_PATH on device before running tests — e.g. "/tmp" for shared libs in /tmp
        timeout: Max seconds for all tests combined (default 300)
    """
    timeout = timeout or DEFAULT_SCRIPT_TIMEOUT
    conn = connection.lower().strip()
    resolved = _resolve_device(device_type, conn)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)
    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    ld = ld_library_path or test_dir
    script_lines = [f"export LD_LIBRARY_PATH={ld}:$LD_LIBRARY_PATH", f"cd {test_dir}"]
    for binary in test_binaries:
        script_lines.append(f'echo "===TEST_START:{binary}==="')
        script_lines.append(f"chmod +x {binary} 2>/dev/null")
        script_lines.append(f"./{binary} 2>&1; echo \"===TEST_EXIT:$?===\"")
        script_lines.append(f'echo "===TEST_END:{binary}==="')
    script = "\n".join(script_lines)

    # WiFi hop if needed
    original_ssid = None
    if conn == "wifi":
        if not device_ssid:
            return json.dumps({"error": "device_ssid required for WiFi."}, indent=2)
        original_ssid = _get_current_ssid()
        ok, msg = _wifi_connect(device_ssid)
        if not ok:
            _wifi_connect(original_ssid)
            return json.dumps({"error": msg}, indent=2)

    t0 = time.time()
    result = _ssh_exec_script(ip, user, pw, script, timeout)

    if original_ssid:
        _wifi_connect(original_ssid)
        _wait_for_internet()

    # Parse GTest results
    raw = result.get("stdout", "")
    test_results = []
    for binary in test_binaries:
        start_marker = f"===TEST_START:{binary}==="
        end_marker = f"===TEST_END:{binary}==="
        si = raw.find(start_marker)
        ei = raw.find(end_marker)
        if si >= 0 and ei >= 0:
            chunk = raw[si + len(start_marker):ei]
            # Extract exit code
            exit_code = -1
            exit_marker = "===TEST_EXIT:"
            eidx = chunk.rfind(exit_marker)
            if eidx >= 0:
                try:
                    exit_code = int(chunk[eidx + len(exit_marker):].split("===")[0])
                except ValueError:
                    pass
                chunk = chunk[:eidx].strip()

            # Parse GTest summary
            passed = failed = total = 0
            for line in chunk.splitlines():
                if line.strip().startswith("[  PASSED  ]"):
                    try:
                        passed = int(line.strip().split("]")[1].strip().split()[0])
                    except (IndexError, ValueError):
                        pass
                elif line.strip().startswith("[  FAILED  ]"):
                    try:
                        failed = int(line.strip().split("]")[1].strip().split()[0])
                    except (IndexError, ValueError):
                        pass

            total = passed + failed
            # Truncate output to avoid huge responses
            output_preview = chunk[-500:] if len(chunk) > 500 else chunk

            test_results.append({
                "binary": binary,
                "exit_code": exit_code,
                "passed": passed,
                "failed": failed,
                "total": total,
                "status": "PASSED" if exit_code == 0 else "FAILED",
                "output_tail": output_preview,
            })
        else:
            test_results.append({
                "binary": binary,
                "status": "NOT_FOUND",
                "output_tail": "",
            })

    total_passed = sum(t.get("passed", 0) for t in test_results)
    total_failed = sum(t.get("failed", 0) for t in test_results)

    return json.dumps({
        "device_type": device_type.upper(),
        "connection": connection,
        "summary": {
            "total_passed": total_passed,
            "total_failed": total_failed,
            "binaries_passed": sum(1 for t in test_results if t.get("status") == "PASSED"),
            "binaries_failed": sum(1 for t in test_results if t.get("status") == "FAILED"),
            "binaries_total": len(test_results),
        },
        "test_results": test_results,
        "duration": f"{time.time() - t0:.1f}s",
    }, indent=2)


# ---------------------------------------------------------------------------
# MEDIUM VALUE: device_logs — fetch logs from device
# ---------------------------------------------------------------------------

@mcp.tool()
def device_logs(device_type: str, log_type: str = "dmesg",
                lines: int = 100, connection: str = "ethernet",
                device_ssid: str = "", device_ip: str = "",
                ssh_user: str = "", ssh_password: str = "",
                grep_filter: str = "") -> str:
    """
    Fetch log output from a device. Supports kernel logs (dmesg), system logs (syslog/journal), or any custom log file.
    Returns the last N lines, optionally filtered by a grep pattern.
    Works over Ethernet, USB, or WiFi (auto-hops for WiFi).

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        log_type: What to fetch — "dmesg" (kernel), "syslog" (system), "journal" (systemd), or a file path like "/var/log/app.log"
        lines: Number of lines to fetch from the end (default 100)
        connection: "ethernet", "usb", or "wifi"
        device_ssid: Required only for WiFi
        device_ip: Override auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
        grep_filter: Optional pattern to filter — e.g. "error" or "USB"
    """
    conn = connection.lower().strip()
    resolved = _resolve_device(device_type, conn)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)
    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    log_cmds = {
        "dmesg": f"dmesg | tail -n {lines}",
        "syslog": f"tail -n {lines} /var/log/syslog 2>/dev/null || tail -n {lines} /var/log/messages 2>/dev/null || echo 'No syslog found'",
        "journal": f"journalctl -n {lines} --no-pager 2>/dev/null || echo 'No journalctl'",
    }
    cmd = log_cmds.get(log_type, f"tail -n {lines} {log_type}")
    if grep_filter:
        cmd += f" | grep -i '{grep_filter}'"

    original_ssid = None
    if conn == "wifi":
        if not device_ssid:
            return json.dumps({"error": "device_ssid required for WiFi."}, indent=2)
        original_ssid = _get_current_ssid()
        ok, msg = _wifi_connect(device_ssid)
        if not ok:
            _wifi_connect(original_ssid)
            return json.dumps({"error": msg}, indent=2)

    results = _ssh_exec(ip, user, pw, [cmd], per_cmd_timeout=30)

    if original_ssid:
        _wifi_connect(original_ssid)
        _wait_for_internet()

    return json.dumps({
        "device_type": device_type.upper(),
        "log_type": log_type,
        "lines_requested": lines,
        "grep_filter": grep_filter,
        "output": results[0] if results else {},
    }, indent=2)


# ---------------------------------------------------------------------------
# MEDIUM VALUE: device_service — manage systemd services
# ---------------------------------------------------------------------------

@mcp.tool()
def device_service(device_type: str, service_name: str,
                   action: str = "status", connection: str = "ethernet",
                   device_ssid: str = "", device_ip: str = "",
                   ssh_user: str = "", ssh_password: str = "") -> str:
    """
    Control a systemd service on a device: check its status, start, stop, restart, enable, disable, or view its logs.
    Does ONE action per call. Works over Ethernet, USB, or WiFi (auto-hops for WiFi).

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        service_name: Name of the systemd service — e.g. "my-app", "nginx", "sshd" (any systemd service name)
        action: What to do — "status", "start", "stop", "restart", "enable", "disable", or "logs"
        connection: "ethernet", "usb", or "wifi"
        device_ssid: Required only for WiFi
        device_ip: Override auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
    """
    valid_actions = {"start", "stop", "restart", "status", "enable", "disable", "logs"}
    if action not in valid_actions:
        return json.dumps({"error": f"Invalid action: {action}. Use: {valid_actions}"}, indent=2)

    conn = connection.lower().strip()
    resolved = _resolve_device(device_type, conn)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)
    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    if action == "logs":
        cmd = f"journalctl -u {service_name} -n 50 --no-pager 2>/dev/null || echo 'No journalctl'"
    else:
        cmd = f"systemctl {action} {service_name} 2>&1; echo '---'; systemctl is-active {service_name} 2>&1"

    original_ssid = None
    if conn == "wifi":
        if not device_ssid:
            return json.dumps({"error": "device_ssid required for WiFi."}, indent=2)
        original_ssid = _get_current_ssid()
        ok, msg = _wifi_connect(device_ssid)
        if not ok:
            _wifi_connect(original_ssid)
            return json.dumps({"error": msg}, indent=2)

    results = _ssh_exec(ip, user, pw, [cmd], per_cmd_timeout=30)

    if original_ssid:
        _wifi_connect(original_ssid)
        _wait_for_internet()

    return json.dumps({
        "device_type": device_type.upper(),
        "service": service_name,
        "action": action,
        "result": results[0] if results else {},
    }, indent=2)


# ---------------------------------------------------------------------------
# MEDIUM VALUE: device_process_list — what's running
# ---------------------------------------------------------------------------

@mcp.tool()
def device_process_list(device_type: str, connection: str = "ethernet",
                        device_ssid: str = "", device_ip: str = "",
                        ssh_user: str = "", ssh_password: str = "",
                        filter_name: str = "") -> str:
    """
    List running processes on a device, sorted by CPU usage (top 30).
    Optionally filter by process name. Works over Ethernet, USB, or WiFi (auto-hops for WiFi).

    Args:
        device_type: Device name — e.g. "DeviceA", "DeviceB" (as configured in .env)
        connection: "ethernet", "usb", or "wifi"
        device_ssid: Required only for WiFi
        device_ip: Override auto-detected IP
        ssh_user: Override SSH username
        ssh_password: Override SSH password
        filter_name: Only show processes matching this name — e.g. "my-app" or "python"
    """
    conn = connection.lower().strip()
    resolved = _resolve_device(device_type, conn)
    if "error" in resolved:
        return json.dumps(resolved, indent=2)
    ip = device_ip or resolved["ip"]
    user = ssh_user or resolved["user"]
    pw = ssh_password or resolved["password"]

    cmd = "ps aux --sort=-%cpu 2>/dev/null || ps -ef"
    if filter_name:
        cmd += f" | grep -i '{filter_name}' | grep -v grep"
    else:
        cmd += " | head -30"

    original_ssid = None
    if conn == "wifi":
        if not device_ssid:
            return json.dumps({"error": "device_ssid required for WiFi."}, indent=2)
        original_ssid = _get_current_ssid()
        ok, msg = _wifi_connect(device_ssid)
        if not ok:
            _wifi_connect(original_ssid)
            return json.dumps({"error": msg}, indent=2)

    results = _ssh_exec(ip, user, pw, [cmd], per_cmd_timeout=15)

    if original_ssid:
        _wifi_connect(original_ssid)
        _wait_for_internet()

    return json.dumps({
        "device_type": device_type.upper(),
        "connection": connection,
        "filter": filter_name,
        "processes": results[0] if results else {},
    }, indent=2)


# ---------------------------------------------------------------------------
# MEDIUM VALUE: multi_device_ssh — run command on multiple devices
# ---------------------------------------------------------------------------

@mcp.tool()
def multi_device_ssh(devices: list[dict], commands: list[str],
                     per_cmd_timeout: int = 0) -> str:
    """
    Run the same shell commands on multiple devices in sequence and collect all results.
    Useful for checking firmware versions, disk space, or running the same test across devices.
    Only for Ethernet/USB (direct connections). For WiFi devices, call wifi_hop_ssh per device.

    Args:
        devices: List of devices, each a dict with "device_type" and "connection" — e.g. [{"device_type": "DeviceA", "connection": "ethernet"}, {"device_type": "DeviceB", "connection": "usb"}]
        commands: Shell commands to run on EACH device — e.g. ["uname -a", "cat /etc/version"]
        per_cmd_timeout: Max seconds per command (default 300)
    """
    per_cmd_timeout = per_cmd_timeout or DEFAULT_CMD_TIMEOUT
    all_results = []

    for dev in devices:
        dtype = dev.get("device_type", "")
        conn = dev.get("connection", "ethernet")
        dev_ip = dev.get("device_ip", "")

        resolved = _resolve_device(dtype, conn)
        if "error" in resolved:
            all_results.append({
                "device_type": dtype,
                "connection": conn,
                "error": resolved["error"],
            })
            continue

        ip = dev_ip or resolved["ip"]
        user = resolved["user"]
        pw = resolved["password"]

        t0 = time.time()
        results = _ssh_exec(ip, user, pw, commands, per_cmd_timeout)

        all_results.append({
            "device_type": dtype.upper(),
            "connection": conn,
            "device_ip": ip,
            "command_results": results,
            "duration": f"{time.time() - t0:.1f}s",
        })

    return json.dumps({"device_results": all_results, "device_count": len(all_results)}, indent=2)


# ---------------------------------------------------------------------------
# WiFi switch / status
# ---------------------------------------------------------------------------

@mcp.tool()
def wifi_switch(ssid: str) -> str:
    """
    Switch WiFi to a different office network (one that has internet). Use this to move between office networks.
    BLOCKED for device SSIDs (as configured in DEVICE_PREFIXES in .env) — those have no internet and would kill the Claude connection. Use wifi_hop_ssh to interact with devices instead.
    """
    if _is_device_ssid(ssid):
        return json.dumps({
            "error": f"BLOCKED: '{ssid}' is a device network (no internet). Use wifi_hop_ssh instead.",
            "example": f'wifi_hop_ssh(device_ssid="{ssid}", commands=["ls /tmp"])',
        }, indent=2)

    current = _get_current_ssid()
    ok, msg = _wifi_connect(ssid)
    return json.dumps({"previous_ssid": current, "target_ssid": ssid, "success": ok, "message": msg}, indent=2)


@mcp.tool()
def wifi_status() -> str:
    """Show which WiFi network this machine is currently connected to, including signal strength, band, channel, and speed."""
    ssid = _get_current_ssid()
    raw = _run_cmd("netsh wlan show interfaces", timeout=10)
    info = {"connected_ssid": ssid}
    for line in raw.splitlines():
        line = line.strip()
        if line.startswith("State"):
            info["state"] = line.split(":", 1)[1].strip()
        elif line.startswith("Band"):
            info["band"] = line.split(":", 1)[1].strip()
        elif line.startswith("Channel"):
            info["channel"] = line.split(":", 1)[1].strip()
        elif line.startswith("Signal"):
            info["signal"] = line.split(":", 1)[1].strip()
        elif line.startswith("Receive rate"):
            info["receive_rate"] = line.split(":", 1)[1].strip()
        elif line.startswith("Transmit rate"):
            info["transmit_rate"] = line.split(":", 1)[1].strip()
    return json.dumps(info, indent=2)


if __name__ == "__main__":
    mcp.run()
