"""Bounded TCP device discovery shared by the console, CLI and MCP.

USB and Android SDK emulators are discovered by ADB itself. Third-party
emulators expose TCP adbd listeners; only listeners owned by known emulator
processes are probed. Explicitly connected TCP devices are remembered separately.
"""

from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path
import socket
import struct
import subprocess
import tempfile
import threading
import time

import psutil

from artemis.config import settings
from artemis.config.paths import get_app_dir
from artemis.runtime.adb_endpoint import AdbSession, current_adb_endpoint
from artemis.utils.logger import get_logger

logger = get_logger(__name__)


def _canonical_target(target: str) -> str:
    host, separator, port = target.rpartition(":")
    if separator and host.lower() in {"localhost", "127.0.0.1"}:
        return f"127.0.0.1:{port}"
    return target


def emulator_listeners() -> list[list[str]]:
    """Return listener groups, one per running emulator (never scan a network)."""
    owners = set()
    for process in psutil.process_iter(["pid", "name", "exe"]):
        info = process.info
        name = str(info.get("name") or "").lower().removesuffix(".exe")
        path = str(info.get("exe") or "").lower()
        known = name in {
            "mumunxdevice",
            "nemuheadless",
            "mumuheadless",
            "ldvboxheadless",
            "noxvmhandle",
            "memuheadless",
            "hd-player",
        }
        vendor_vm = name == "vboxheadless" and any(
            vendor in path for vendor in ("mumu", "nemu", "ldplayer", "leidian", "nox", "memu")
        )
        if known or vendor_vm:
            owners.add(info["pid"])
    if not owners:
        return []
    groups: dict[int, set[int]] = {}
    for connection in psutil.net_connections(kind="tcp"):
        address = connection.laddr
        pid = connection.pid
        if (
            pid is not None
            and pid in owners
            and address
            and connection.status == psutil.CONN_LISTEN
            and address.ip in {"127.0.0.1", "0.0.0.0"}
        ):
            groups.setdefault(pid, set()).add(address.port)
    # Prefer current MuMu ports over its compatibility aliases (same adbd).
    return [
        [f"127.0.0.1:{port}" for port in sorted(ports, key=lambda p: (p in {5555, 7555}, p))]
        for ports in groups.values()
    ]


def is_adbd(target: str, timeout: float) -> bool:
    """Confirm an ADB CNXN/AUTH header before giving a local port to adb connect."""
    host, port = target.rsplit(":", 1)
    payload = b"host::\x00"
    command = int.from_bytes(b"CNXN", "little")
    packet = struct.pack(
        "<6I", command, 0x01000000, 4096, len(payload), sum(payload), command ^ 0xFFFFFFFF
    )
    try:
        with socket.create_connection((host, int(port)), timeout=timeout) as stream:
            stream.sendall(packet + payload)
            header = b""
            deadline = time.monotonic() + timeout
            while len(header) < 24:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    return False
                stream.settimeout(remaining)
                chunk = stream.recv(24 - len(header))
                if not chunk:
                    return False
                header += chunk
            kind, _, _, length, _, magic = struct.unpack("<6I", header)
            return (
                kind.to_bytes(4, "little") in {b"CNXN", b"AUTH"}
                and magic == kind ^ 0xFFFFFFFF
                and length <= 1048576
            )
    except (OSError, ValueError):
        return False


class AdbDeviceDiscovery:
    """Coalesce discovery and reconnect only missing transports, with a total budget."""

    def __init__(self):
        self._gate = threading.Lock()
        self._last_scan = float("-inf")
        self.results: list[dict[str, object]] = []
        self._cursor = 0

    @staticmethod
    def _directory() -> Path:
        return get_app_dir() / "adb-devices"

    def remember(self, target: str) -> bool:
        """Save an explicitly connected device, without credentials or shared-list races."""
        directory = self._directory()
        temporary = None
        try:
            directory.mkdir(parents=True, exist_ok=True)
            destination = directory / (hashlib.sha256(target.encode()).hexdigest() + ".json")
            fd, temporary = tempfile.mkstemp(dir=directory, prefix=".device-")
            with os.fdopen(fd, "w", encoding="utf-8") as stream:
                json.dump({"target": target}, stream)
            os.replace(temporary, destination)
            return True
        except OSError as exc:
            logger.warning(f"Cannot remember ADB device {target}: {exc}")
            return False
        finally:
            if temporary:
                Path(temporary).unlink(missing_ok=True)

    def _remembered(self) -> list[str]:
        targets = []
        for path in self._directory().glob("*.json"):
            try:
                target = json.loads(path.read_text(encoding="utf-8"))["target"]
                host, port = target.rsplit(":", 1)
                if host and 0 < int(port) < 65536:
                    targets.append(target)
            except (OSError, ValueError, KeyError, TypeError, AttributeError) as exc:
                logger.warning(f"Invalid remembered ADB device {path.name}: {exc}")
        return targets

    def refresh(self, adb_path: str, *, force: bool = False) -> list[dict[str, object]]:
        if not settings.ARTEMIS_ADB_AUTO_CONNECT or not current_adb_endpoint().is_local_default:
            self.results = []
            return []
        if not self._gate.acquire(blocking=False):
            return self.results.copy()
        try:
            if (
                not force
                and time.monotonic() - self._last_scan
                < settings.ARTEMIS_ADB_DISCOVERY_INTERVAL_SECONDS
            ):
                return self.results.copy()
            self._last_scan = time.monotonic()
            self.results = []
            self._refresh(adb_path)
        except (OSError, psutil.Error, ValueError, subprocess.SubprocessError) as exc:
            logger.warning(f"ADB automatic discovery unavailable: {exc}")
            self.results.append({"success": False, "message": str(exc)})
        finally:
            self._gate.release()
        return self.results.copy()

    def _refresh(self, adb_path: str) -> None:
        session = AdbSession(current_adb_endpoint(), adb_path=adb_path)
        deadline = time.monotonic() + settings.ARTEMIS_ADB_DISCOVERY_TIMEOUT_SECONDS
        devices = session.run(
            ["devices"],
            capture_output=True,
            text=True,
            timeout=max(deadline - time.monotonic(), 0.01),
        )
        if devices.returncode:
            return
        ready = {
            _canonical_target(parts[0])
            for line in devices.stdout.splitlines()
            if len(parts := line.split()) >= 2 and parts[1] in {"device", "unauthorized"}
        }
        try:
            groups = emulator_listeners()
        except (OSError, psutil.Error) as exc:
            # Missing process-inspection permissions must not stop reconnection
            # of explicitly configured phones or make device enumeration fail.
            groups = []
            logger.warning(f"Cannot inspect local emulator listeners: {exc}")
            self.results.append({"success": False, "message": str(exc)})
        remembered = self._remembered()
        targets = [(group, True) for group in groups] + [([target], False) for target in remembered]
        # Resume at a different group if a slow/offline peer exhausts the budget.
        # Large multi-instance installations must not starve later devices.
        count = len(targets)
        start = self._cursor % count if count else 0
        for offset in range(count):
            index = (start + offset) % count
            candidates, probe = targets[index]
            self._cursor = (index + 1) % count
            if any(_canonical_target(target) in ready for target in candidates):
                continue
            for target in candidates:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    return
                if probe and not is_adbd(target, min(remaining, 0.3)):
                    continue
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    return
                try:
                    result = session.run(
                        ["connect", target],
                        capture_output=True,
                        text=True,
                        timeout=min(remaining, 2),
                    )
                except subprocess.TimeoutExpired:
                    self.results.append(
                        {"target": target, "success": False, "message": "ADB 连接超时"}
                    )
                    continue
                output = (result.stdout + result.stderr).strip()
                success = (
                    result.returncode == 0
                    and "connected to" in output.lower()
                    and "failed" not in output.lower()
                )
                state = self._wait_for_transport(session, target, deadline) if success else None
                pending = success and state is None
                if success and state is None:
                    success = False
                    output = f"{output}；已发起连接，等待设备就绪，下次刷新会重试"
                self.results.append(
                    {
                        "target": target,
                        "success": success,
                        "state": state,
                        "pending": pending,
                        "message": output,
                    }
                )
                if success:
                    logger.info(f"ADB 自动连接：{target} — {output}")
                    ready.add(_canonical_target(target))
                    break
                if pending:
                    logger.info(f"ADB 设备连接中：{target} — {output}")
                    break  # Do not create a second alias while the first transport is settling.
                logger.warning(f"ADB 自动连接失败：{target} — {output}")

    @staticmethod
    def _wait_for_transport(session: AdbSession, target: str, deadline: float) -> str | None:
        """adb connect can return before its transport appears in adb devices."""
        while (remaining := deadline - time.monotonic()) > 0:
            try:
                result = session.run(["devices"], capture_output=True, text=True, timeout=remaining)
            except subprocess.TimeoutExpired:
                return None
            for line in result.stdout.splitlines():
                parts = line.split()
                if (
                    len(parts) >= 2
                    and _canonical_target(parts[0]) == _canonical_target(target)
                    and parts[1] in {"device", "unauthorized"}
                ):
                    return parts[1]
            time.sleep(min(0.1, max(deadline - time.monotonic(), 0)))
        return None


adb_device_discovery = AdbDeviceDiscovery()