"""
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 from same directory as this script
# ---------------------------------------------------------------------------
_SCRIPT_DIR = Path(__file__).resolve().parent
_ENV_FILE = _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} e.g. {"CLM": "CEILING"}"""
    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 'CLM', look for 'CEILING_xxx' SSIDs."
    ),
)


# ---------------------------------------------------------------------------
# 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.
    e.g. CEILING_D0A0117F2A11 → checks CEILING in profiles,
    and also checks if any alias points to CEILING."""
    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:
    """Scan for available WiFi networks."""
    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 with status, MAC, speed, and IP config."""
    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 connected USB devices."""
    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:
    """Scan all WiFi, network adapters, and USB devices in one call."""
    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. CLM -> CEILING)
    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 SSH commands on a device via Ethernet or USB (no WiFi hop needed).
    Internet stays connected — these are direct wired connections.

    Args:
        device_type: Device name (e.g. "STPRO", "CLM", "CEILING", "MLM2")
        commands: List of shell commands to run
        connection: "ethernet" or "usb" (default: ethernet)
        device_ip: Override device IP (auto-detected from device type + connection)
        ssh_user: Override SSH user
        ssh_password: Override SSH password
        per_cmd_timeout: Timeout per command in seconds
    """
    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 full bash script on a device via Ethernet or USB (no WiFi hop needed).

    Args:
        device_type: Device name (e.g. "STPRO", "CLM", "CEILING")
        script: Full bash script to execute
        connection: "ethernet" or "usb"
        device_ip: Override device IP
        ssh_user: Override SSH user
        ssh_password: Override SSH password
        timeout: Script timeout in seconds
    """
    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:
    """
    Push a file to a device via Ethernet or USB (no WiFi hop needed).

    Args:
        device_type: Device name (e.g. "STPRO", "CLM", "CEILING")
        local_path: Local file path to send
        remote_path: Destination path on device
        connection: "ethernet" or "usb"
        device_ip: Override device IP
        ssh_user: Override SSH user
        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:
    """
    Pull a file from a device via Ethernet or USB (no WiFi hop needed).

    Args:
        device_type: Device name (e.g. "STPRO", "CLM", "CEILING")
        remote_path: File path on device to download
        local_path: Local destination path
        connection: "ethernet" or "usb"
        device_ip: Override device IP
        ssh_user: Override SSH user
        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:
    """
    Hop WiFi to a device, push a file via SCP/SFTP, then hop back.

    Args:
        device_ssid: Device WiFi SSID
        local_path: Local file path to send (Windows path)
        remote_path: Destination path on device
        home_ssid: Office WiFi to return to (auto-detected)
        device_ip: Device IP (auto-detected)
        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:
    """
    Hop WiFi to a device, pull a file via SCP/SFTP, then hop back.

    Args:
        device_ssid: Device WiFi SSID
        remote_path: File path on device to download
        local_path: Local destination path (Windows path)
        home_ssid: Office WiFi to return to (auto-detected)
        device_ip: Device IP (auto-detected)
        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)

    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:
    """
    Hop WiFi to a device, run SSH commands, then hop back to office WiFi.

    1. Remembers current office WiFi
    2. Switches to device WiFi
    3. Runs SSH commands on the device
    4. Switches back to office WiFi
    5. Verifies internet is working

    Args:
        device_ssid: Device WiFi SSID (e.g. "MYDEVICE_ABC123")
        commands: List of shell commands to run on device
        home_ssid: Office WiFi to return to (auto-detected if empty)
        device_ip: Device IP (auto-detected from device type)
        ssh_user: SSH user (auto-detected from device type)
        ssh_password: SSH password (auto-detected from device type)
        per_cmd_timeout: Timeout per command in seconds (default from .env)
    """
    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:
    """
    Hop WiFi to a device, run a full bash script, then hop back.

    Args:
        device_ssid: Device WiFi SSID
        script: Full bash script to run on device
        home_ssid: Office WiFi to return to
        device_ip: Device IP
        ssh_user: SSH user
        ssh_password: SSH password
        timeout: Script timeout in seconds (default from .env)
    """
    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)


# ---------------------------------------------------------------------------
# HIGH VALUE: device_deploy — tar, push, extract in one call
# ---------------------------------------------------------------------------

@mcp.tool()
def device_deploy(device_type: str, local_dir: str, remote_dir: str = "/tmp",
                  connection: str = "ethernet", device_ssid: str = "",
                  file_glob: str = "", device_ip: str = "",
                  ssh_user: str = "", ssh_password: str = "") -> str:
    """
    Deploy files from a local directory to a device in one step.
    Tars the local directory, pushes via SFTP, and extracts on device.

    Works with all connections: WiFi (auto-hops), Ethernet, USB.

    Args:
        device_type: Device name (e.g. "CLM", "STPRO")
        local_dir: Local directory to deploy (will be tarred)
        remote_dir: Destination directory on device (default /tmp)
        connection: "wifi", "ethernet", or "usb"
        device_ssid: Required for WiFi — the device SSID
        file_glob: Optional glob to filter files (e.g. "*.so" or "test_*")
        device_ip: Override device IP
        ssh_user: Override SSH user
        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)

    # Resolve device
    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()

    # Create tarball
    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 ({time.time() - t0:.1f}s)")
    except Exception as e:
        return json.dumps({"error": f"Tar failed: {e}"}, indent=2)

    # WiFi hop or direct
    original_ssid = None
    if conn == "wifi":
        if not device_ssid:
            os.unlink(tmp_tar)
            return json.dumps({"error": "device_ssid required for WiFi connection."}, 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)
        log.append(f"WiFi hop: {time.time() - t0:.1f}s")

    # Push tarball
    t1 = time.time()
    remote_tar = f"{remote_dir}/deploy_package.tar.gz"
    push_result = _sftp_push(ip, user, pw, tmp_tar, remote_tar)
    os.unlink(tmp_tar)
    if not push_result.get("success"):
        if original_ssid:
            _wifi_connect(original_ssid)
        return json.dumps({"error": f"Push failed: {push_result.get('error')}", "log": log}, indent=2)
    log.append(f"SFTP push: {time.time() - t1:.1f}s")

    # Extract on device
    t2 = time.time()
    extract_result = _ssh_exec(ip, user, pw, [
        f"cd {remote_dir} && tar xzf deploy_package.tar.gz && rm deploy_package.tar.gz && ls -la"
    ])
    log.append(f"Extract: {time.time() - t2:.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,
        "bytes_transferred": tar_size,
        "remote_dir": remote_dir,
        "extract_result": extract_result[0] if extract_result else {},
        "home_restored": home_restored,
        "log": log,
    }, 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:
    """
    Quick health check: CPU, RAM, disk, uptime, temperature, running processes.
    Works with WiFi (auto-hops), Ethernet, or USB.

    Args:
        device_type: Device name (e.g. "CLM", "STPRO")
        connection: "wifi", "ethernet", or "usb"
        device_ssid: Required for WiFi
        device_ip: Override IP
        ssh_user: Override user
        ssh_password: Override 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:
    """
    Check if a device is reachable (ping test). Use before SSH to avoid timeouts.
    No WiFi hop — only works for Ethernet/USB or already-connected WiFi.

    Args:
        device_type: Device name (e.g. "CLM", "STPRO")
        connection: "wifi", "ethernet", or "usb"
        device_ip: Override 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 for it to come back online.
    Works with WiFi (auto-hops), Ethernet, or USB.

    Args:
        device_type: Device name (e.g. "CLM", "STPRO")
        connection: "wifi", "ethernet", or "usb"
        device_ssid: Required for WiFi
        device_ip: Override IP
        ssh_user: Override user
        ssh_password: Override password
        wait_timeout: Seconds to wait for device to come back (0 = 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 test binaries on device and return structured pass/fail results.
    Parses GTest output automatically.

    Args:
        device_type: Device name (e.g. "CLM", "STPRO")
        test_binaries: List of test binary names (e.g. ["smoke_test", "test_iox_wrapper"])
        test_dir: Directory on device where binaries are (default /tmp)
        connection: "wifi", "ethernet", or "usb"
        device_ssid: Required for WiFi
        device_ip: Override IP
        ssh_user: Override user
        ssh_password: Override password
        ld_library_path: Extra LD_LIBRARY_PATH on device (e.g. "/tmp")
        timeout: Total timeout for all tests
    """
    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 logs from a device. Supports dmesg, syslog, journalctl, or custom log files.

    Args:
        device_type: Device name (e.g. "CLM", "STPRO")
        log_type: "dmesg", "syslog", "journal", or a file path (e.g. "/var/log/app.log")
        lines: Number of lines to fetch (default 100)
        connection: "wifi", "ethernet", or "usb"
        device_ssid: Required for WiFi
        device_ip: Override IP
        ssh_user: Override user
        ssh_password: Override password
        grep_filter: Optional grep pattern to filter logs
    """
    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:
    """
    Manage a systemd service on device: start, stop, restart, status, enable, disable.

    Args:
        device_type: Device name (e.g. "CLM", "STPRO")
        service_name: Service name (e.g. "DeviceMain", "nginx")
        action: "start", "stop", "restart", "status", "enable", "disable", "logs"
        connection: "wifi", "ethernet", or "usb"
        device_ssid: Required for WiFi
        device_ip: Override IP
        ssh_user: Override user
        ssh_password: Override 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 device, sorted by CPU usage.

    Args:
        device_type: Device name (e.g. "CLM", "STPRO")
        connection: "wifi", "ethernet", or "usb"
        device_ssid: Required for WiFi
        device_ip: Override IP
        ssh_user: Override user
        ssh_password: Override password
        filter_name: Optional process name filter (grep)
    """
    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 SSH commands on multiple devices sequentially.
    Only for Ethernet/USB (direct connections). For WiFi, use wifi_hop_ssh per device.

    Args:
        devices: List of {device_type, connection} dicts.
                 e.g. [{"device_type": "STPRO", "connection": "ethernet"},
                       {"device_type": "CLM", "connection": "usb"}]
        commands: List of shell commands to run on each device
        per_cmd_timeout: Timeout per command
    """
    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 an office/internet network ONLY.
    BLOCKED for device SSIDs — use wifi_hop_ssh 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:
    """Get current WiFi connection status."""
    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()
