#!/usr/bin/env python3
"""USB 热拔插监控库（Windows）- 基于 cyme --json watch

依赖 cyme.exe，可通过 download_cyme() 函数自动下载，或通过 cyme_path 参数手动指定。

使用示例:
    import cyme_win

    # 方式一：简单回调（推荐）
    cyme_win.watch(lambda d: print(d['name'], d['event']))

    # 方式二：指定 cyme.exe 路径
    cyme_win.watch(lambda d: print(d), cyme_path="E:/tools/cyme.exe")

    # 方式三：自定义监控器
    monitor = cyme_win.USBMonitor(cyme_path="E:/tools/cyme.exe")
    monitor.add_callback(my_handler).start()

设备字典字段:
    - event: 事件类型 ("connect" 或 "disconnect")
    - name: 设备名称
    - vendor_id: 厂商ID (如 "046d")
    - product_id: 产品ID (如 "c52b")
    - vid_pid: VID:PID 格式 (如 "046d:c52b")
    - bus: 总线号
    - port: 端口号
    - location: 位置 (如 "1-2")
    - manufacturer: 厂商名称
    - serial: 序列号
    - speed: USB 速度
    - device_class: 设备类
    - raw_event: 原始事件数据
"""
from __future__ import annotations

import asyncio
import json
import os
import shutil
import sys
import threading
import urllib.request
import zipfile
from typing import Callable

# cyme.exe 路径缓存（延迟查找，避免导入时报错）
_cyme_path: str | None = None


def download_cyme(target_dir: str | None = None) -> str:
    """下载 cyme.exe 到指定目录。

    从 GitHub Releases 获取最新版本的 cyme.exe。
    如果目标路径已存在，则直接返回该路径。

    Args:
        target_dir: 目标目录，默认为当前工作目录。

    Returns:
        cyme.exe 的完整路径。

    Raises:
        RuntimeError: 网络错误或未找到下载链接。

    Example:
        >>> cyme_win.download_cyme()  # 下载到当前目录
        >>> cyme_win.download_cyme("E:/my_project")  # 下载到指定目录
    """
    directory = target_dir or os.getcwd()
    target_path = os.path.join(directory, "cyme.exe")

    # 已存在则直接返回
    if os.path.isfile(target_path):
        print(f"cyme.exe 已存在: {target_path}")
        return target_path

    # 从 GitHub API 获取最新版本信息
    api = "https://api.github.com/repos/tuna-f1sh/cyme/releases/latest"
    print("正在获取 cyme 最新版本...")
    with urllib.request.urlopen(api, timeout=10) as resp:
        data = json.loads(resp.read().decode())

    version = data["tag_name"]
    url = None
    for asset in data.get("assets", []):
        name = asset["name"]
        if "windows" in name.lower() and name.endswith(".zip"):
            url = asset["browser_download_url"]
            break

    if not url:
        raise RuntimeError("未找到 Windows 版本下载链接")

    # 下载并解压
    print(f"下载 cyme {version}...")
    zip_path = os.path.join(directory, "cyme.zip")
    urllib.request.urlretrieve(url, zip_path)

    with zipfile.ZipFile(zip_path, "r") as z:
        for name in z.namelist():
            if name.endswith("cyme.exe"):
                with z.open(name) as src, open(target_path, "wb") as dst:
                    dst.write(src.read())
                break

    os.remove(zip_path)
    print(f"cyme {version} 已下载到: {target_path}")
    return target_path


def _get_cyme_path(cyme_path: str | None = None) -> str:
    """获取 cyme.exe 路径。

    Args:
        cyme_path: 手动指定的路径，如果提供则直接使用。

    查找顺序：
        1. 手动指定的路径
        2. 当前工作目录
        3. 可执行文件所在目录（PyInstaller 打包后）
        4. 模块所在目录（pip 安装）
        5. 系统 PATH

    Returns:
        cyme.exe 的完整路径。

    Raises:
        RuntimeError: 未找到 cyme.exe。
    """
    global _cyme_path

    # 1. 手动指定的路径
    if cyme_path:
        if os.path.isfile(cyme_path):
            return cyme_path
        raise RuntimeError(f"指定的 cyme.exe 不存在: {cyme_path}")

    # 使用缓存
    if _cyme_path:
        return _cyme_path

    # 2. 当前工作目录
    cwd_path = os.path.join(os.getcwd(), "cyme.exe")
    if os.path.isfile(cwd_path):
        _cyme_path = cwd_path
        return _cyme_path

    # 3. 可执行文件所在目录（PyInstaller 打包后）
    if getattr(sys, "frozen", False):
        exe_path = os.path.join(os.path.dirname(sys.executable), "cyme.exe")
        if os.path.isfile(exe_path):
            _cyme_path = exe_path
            return _cyme_path

    # 4. 模块所在目录（pip 安装）
    mod_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cyme.exe")
    if os.path.isfile(mod_path):
        _cyme_path = mod_path
        return _cyme_path

    # 5. 系统 PATH
    path = shutil.which("cyme") or shutil.which("cyme.exe")
    if path:
        _cyme_path = path
        return _cyme_path

    raise RuntimeError(
        "未找到 cyme.exe\n"
        "请先调用 cyme_win.download_cyme() 下载，或手动放到以下位置之一：\n"
        f"  - 当前目录: {os.getcwd()}\n"
        f"  - 模块目录: {os.path.dirname(os.path.abspath(__file__))}\n"
        "下载地址: https://github.com/tuna-f1sh/cyme/releases"
    )


def _make_device_dict(d: dict, bus: int, event_type: str) -> dict:
    """创建设备信息字典。

    将 cyme 输出的设备数据转换为更易用的字典格式。

    Args:
        d: cyme 输出的设备数据。
        bus: 总线号。
        event_type: 事件类型 ("connect" 或 "disconnect")。

    Returns:
        包含设备信息的字典。
    """
    loc = d.get("location_id", {})
    port = loc.get("number", 0)
    vid = d.get("vendor_id", 0)
    pid = d.get("product_id", 0)

    return {
        "event": event_type,
        "name": d.get("name", "Unknown Device"),
        "vendor_id": f"{vid:04x}",
        "product_id": f"{pid:04x}",
        "vid_pid": f"{vid:04x}:{pid:04x}",
        "bus": bus,
        "port": port,
        "location": f"{bus}-{port}",
        "manufacturer": d.get("manufacturer"),
        "serial": d.get("serial_num"),
        "speed": d.get("device_speed"),
        "device_class": d.get("class"),
        "subclass": d.get("sub_class"),
        "protocol": d.get("protocol"),
        "version": d.get("bcd_device"),
        "raw_event": d.get("last_event", {}),
    }


class USBMonitor:
    """USB 热拔插监控器。

    支持链式调用和上下文管理器。

    Example:
        # 链式调用
        monitor = USBMonitor().add_callback(handler).start()

        # 指定 cyme.exe 路径
        monitor = USBMonitor(cyme_path="E:/tools/cyme.exe").add_callback(handler).start()

        # 上下文管理器
        with USBMonitor().add_callback(handler):
            do_something()
    """

    MAX_CACHE = 1000  # 事件缓存上限，防止内存溢出

    def __init__(self, cyme_path: str | None = None):
        """初始化监控器。

        Args:
            cyme_path: cyme.exe 路径，不指定则自动查找。
        """
        self._cyme_path = cyme_path  # cyme.exe 路径
        self._callbacks: list[Callable[[dict], None]] = []  # 回调函数列表
        self._loop: asyncio.AbstractEventLoop | None = None  # asyncio 事件循环
        self._task: asyncio.Task | None = None  # 监控任务
        self._running = False  # 运行状态
        self._seen: set[str] = set()  # 已处理事件ID（用于去重）

    def add_callback(self, callback: Callable[[dict], None]) -> "USBMonitor":
        """添加事件回调函数。

        当 USB 设备插入或拔出时，回调函数会被调用，参数为设备信息字典。

        Args:
            callback: 回调函数，接收一个 dict 参数。

        Returns:
            self，支持链式调用。

        Example:
            monitor.add_callback(lambda d: print(d['name']))
        """
        self._callbacks.append(callback)
        return self

    def start(self) -> "USBMonitor":
        """开始监控。

        在后台线程中启动监控，不会阻塞主线程。

        Returns:
            self，支持链式调用。
        """
        if self._running:
            return self
        self._running = True
        threading.Thread(target=self._run_loop, daemon=True).start()
        return self

    def stop(self) -> None:
        """停止监控。

        终止后台监控线程和 cyme 子进程。
        """
        if not self._running:
            return
        self._running = False
        if self._loop and self._task:
            self._loop.call_soon_threadsafe(self._task.cancel)

    def _run_loop(self) -> None:
        """运行 asyncio 事件循环（内部方法）。

        在独立线程中创建并运行事件循环。
        """
        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)
        self._task = self._loop.create_task(self._monitor())
        try:
            self._loop.run_until_complete(self._task)
        except asyncio.CancelledError:
            pass
        finally:
            self._loop.run_until_complete(self._loop.shutdown_asyncgens())
            self._loop.close()

    async def _monitor(self) -> None:
        """监控 cyme --json watch 输出（内部方法）。

        启动 cyme 子进程，实时读取 JSON 输出并处理。
        """
        proc = await asyncio.create_subprocess_exec(
            _get_cyme_path(self._cyme_path), "--json", "watch",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.DEVNULL,
        )
        try:
            async for line in proc.stdout or ():
                if not self._running:
                    break
                try:
                    self._handle_data(json.loads(line))
                except json.JSONDecodeError:
                    pass
        except asyncio.CancelledError:
            pass
        finally:
            # 确保子进程被终止
            if proc.returncode is None:
                proc.terminate()
                try:
                    await asyncio.wait_for(proc.wait(), timeout=1.0)
                except (asyncio.TimeoutError, asyncio.CancelledError):
                    proc.kill()

    def _handle_data(self, data: dict) -> None:
        """处理 cyme 输出的 JSON 数据（内部方法）。

        解析设备事件，去重后分发给回调函数。

        Args:
            data: cyme 输出的 JSON 数据。
        """
        for bus in data.get("buses", []):
            bus_num = bus.get("usb_bus_number", 0)
            for d in bus.get("devices", []):
                evt = d.get("last_event", {})
                loc = d.get("location_id", {})
                port = loc.get("number", 0)

                # 处理连接事件
                if "connected" in evt:
                    evt_id = f"connect:{d.get('vendor_id')}:{d.get('product_id')}:{bus_num}:{port}:{evt['connected']}"
                    if evt_id not in self._seen:
                        self._seen.add(evt_id)
                        if len(self._seen) > self.MAX_CACHE:
                            self._seen.clear()
                        self._dispatch(d, bus_num, "connect")

                # 处理断开事件
                elif "disconnected" in evt:
                    evt_id = f"disconnect:{d.get('vendor_id')}:{d.get('product_id')}:{bus_num}:{port}:{evt['disconnected']}"
                    if evt_id not in self._seen:
                        self._seen.add(evt_id)
                        if len(self._seen) > self.MAX_CACHE:
                            self._seen.clear()
                        self._dispatch(d, bus_num, "disconnect")

    def _dispatch(self, d: dict, bus: int, event_type: str) -> None:
        """分发事件到回调函数（内部方法）。

        Args:
            d: 设备数据。
            bus: 总线号。
            event_type: 事件类型。
        """
        dev = _make_device_dict(d, bus, event_type)
        for cb in self._callbacks:
            try:
                cb(dev)
            except Exception:
                pass  # 忽略回调异常，避免影响其他回调

    def __enter__(self) -> "USBMonitor":
        """进入上下文，自动启动监控。"""
        return self.start()

    def __exit__(self, *args) -> None:
        """退出上下文，自动停止监控。"""
        self.stop()


# 全局单例监控器
_monitor: USBMonitor | None = None


def watch(callback: Callable[[dict], None], cyme_path: str | None = None) -> USBMonitor:
    """注册 USB 事件回调（全局单例，最简单的调用方式）。

    使用全局单例监控器，适合简单场景。
    如需多个独立监控器，请使用 USBMonitor 类。

    Args:
        callback: 回调函数，接收设备信息字典。
        cyme_path: cyme.exe 路径，不指定则自动查找。

    Returns:
        全局 USBMonitor 实例。

    Example:
        >>> cyme_win.watch(lambda d: print(f"{d['event']}: {d['name']}"))
        >>> cyme_win.watch(lambda d: print(d), cyme_path="E:/tools/cyme.exe")
    """
    global _monitor
    if _monitor is None:
        _monitor = USBMonitor(cyme_path=cyme_path).start()
    return _monitor.add_callback(callback)


def stop() -> None:
    """停止全局监控。

    停止由 watch() 启动的全局监控器。
    """
    global _monitor
    if _monitor:
        _monitor.stop()
        _monitor = None


if __name__ == "__main__":
    import time

    def on_event(device: dict):
        """示例回调函数"""
        symbol = "+" if device["event"] == "connect" else "-"
        print(f"[{symbol}] {device['name']} ({device['vid_pid']})")
        print(f"  位置: {device['location']}")
        print(f"  厂商ID: {device['vendor_id']}  产品ID: {device['product_id']}")
        print(f"  厂商: {device['manufacturer'] or '-'}")
        print(f"  序列号: {device['serial'] or '-'}")
        print(f"  速度: {device['speed'] or '-'}")
        if device.get("device_class"):
            print(f"  设备类: {device['device_class']}")
        print()

    print("USB 监控中... (Ctrl+C 退出)\n")

    monitor = USBMonitor().add_callback(on_event).start()
    try:
        while monitor._running:
            time.sleep(0.5)
    except KeyboardInterrupt:
        print("\n已停止")
    finally:
        monitor.stop()