"""设备生命周期 + 通用 Worker 基类 + 全局状态注册表 + 心跳看门狗。

  STFDevice    — 单设备生命周期：acquire（占用+adb连接）→ release（断开+释放）
  BaseWorker   — 通用 worker 基类：封装 STF 生命周期 + u2 连接 + 状态上报 + stop
                 + 生命周期钩子(setup/teardown/on_error) + 心跳上报
                 子类只需实现 run_task(d)，可选实现钩子，不用管 STF/adb/状态/异常
  _WORKERS     — 全局状态注册表，供 web_server 读取实时状态
  _watchdog    — 心跳看门狗线程，检测 worker 卡死

新增任务只需写一个 BaseWorker 子类 + 一个 BaseTask 子类，
完全不用碰这里的代码。
"""
import time
import threading
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout

import uiautomator2 as u2

from core.logger import get_logger
from core.stf_client import DeviceOfflineError, STFError
from core import device_pool
from .adb_helper import adb_connect, adb_disconnect

_log = get_logger("core.worker")

# 心跳超时阈值（秒）。worker 超过这个时间没更新心跳，判定为卡死
_HEARTBEAT_TIMEOUT = 120
# 看门狗检查间隔
_WATCHDOG_INTERVAL = 30
# u2.connect / d.info 超时（秒）。超过则放弃，避免 worker 永久 hang
_U2_CONNECT_TIMEOUT = 30
_U2_INFO_TIMEOUT = 10


class STFDevice:
    """单设备生命周期：acquire（直连）→ release（无操作）。

    阶段 1：已摘除 STF occupy/release——单实例互斥由 TaskManager._running 保证，
    设备健康由 device_pool（本地 adb 在线状态）判断。保留 STF 引用仅为
    非 IP:port（USB 序列号）设备的 remoteConnect 桥接，阶段 2 换成 220 adb server。
    """

    def __init__(self, stf_client=None, serial=None):
        self.stf = stf_client  # 仅 USB 桥接用（阶段 2 移除）
        self.serial = serial
        self.remote_adb_url = None

    def acquire(self, serial=None):
        self.serial = serial or self.serial or self._pick_free()
        _log.info(f"[{self.serial}] 选定设备")

        # 直连优先：serial 本身是 IP:5555（如 Tailscale 网络设备）时直接 adb connect。
        # 只 connect、绝不 disconnect/kill-server（红线），不影响 STF provider 的连接。
        # 仅当 serial 不是 IP:port（如 USB 序列号）时才走 STF 桥接（阶段 2 替换）。
        if ":" in self.serial:
            self.remote_adb_url = self.serial
            _log.info(f"[{self.serial}] 直连设备: {self.remote_adb_url}")
        else:
            if self.stf is None:
                raise STFError(f"USB 设备 {self.serial} 暂不支持（未启用 STF 桥接）", "offline")
            self.remote_adb_url = self.stf.remote_connect(self.serial)
            _log.info(f"[{self.serial}] STF 桥接: {self.remote_adb_url}")

        if not adb_connect(self.remote_adb_url):
            raise STFError(f"adb connect 失败: {self.remote_adb_url}", "offline")
        time.sleep(2)

    def _pick_free(self):
        free = device_pool.list_ready()
        if not free:
            raise STFError("没有可用的在线空闲设备", "offline")
        return free[0]

    def release(self):
        if not self.serial:
            return
        # 直连模式（remote_adb_url == serial）：不 disconnect（红线，STF 共享 transport）
        if self.remote_adb_url and self.remote_adb_url != self.serial:
            adb_disconnect(self.remote_adb_url)
            try:
                self.stf.remote_disconnect(self.serial)
            except Exception:
                pass


# ================== 全局 worker 状态注册表（供 web_server 读取） ==================
_WORKERS_LOCK = threading.Lock()
_WORKERS = {}  # serial -> status dict


def get_all_worker_status():
    with _WORKERS_LOCK:
        return list(_WORKERS.values())


def _update_status(serial, **fields):
    with _WORKERS_LOCK:
        s = _WORKERS.setdefault(serial, {"serial": serial})
        s.update(fields)
        s["last_heartbeat"] = time.time()


def _remove_worker(serial):
    with _WORKERS_LOCK:
        _WORKERS.pop(serial, None)


def clear_worker_error(serial):
    """清除设备错误/失败状态，恢复为 idle。运行中/连接中返回 False 不清理。

    检查与更新在同一把锁内完成，避免与任务启动写 connecting/running 的竞态。
    """
    with _WORKERS_LOCK:
        w = _WORKERS.get(serial)
        if w is None:
            return True  # 本就没有状态记录，视作已清理
        if w.get("status") in ("running", "connecting"):
            return False
        w.update(status="idle", last_error="", last_warning="", task_job="",
                 current_action="", attempt=0, progress={}, end_time=0,
                 last_heartbeat=time.time())
    return True


def clear_all_worker_errors(exclude=()):
    """批量清除所有 error/failed 状态（exclude 里的跳过）。返回清除数。"""
    exclude = set(exclude)
    with _WORKERS_LOCK:
        targets = [s for s, w in _WORKERS.items()
                   if s not in exclude and w.get("status") in ("error", "failed")]
    cleared = 0
    for s in targets:
        if clear_worker_error(s):
            cleared += 1
    return cleared


def get_worker_heartbeat(serial):
    """返回 worker 最近一次心跳时间，没有返回 0。"""
    with _WORKERS_LOCK:
        return _WORKERS.get(serial, {}).get("last_heartbeat", 0)


# ================== 心跳看门狗 ==================
class _Watchdog(threading.Thread):
    """后台线程：检测 worker 心跳超时，自动标记异常并设置 stop 信号。

    worker 卡死（如 u2 操作 hang 住）时，看门狗能发现并标记状态，
    避免设备一直被占用却不干活。
    """

    def __init__(self):
        super().__init__(daemon=True, name="worker-watchdog")
        self._stop = threading.Event()

    def run(self):
        while not self._stop.wait(_WATCHDOG_INTERVAL):
            try:
                self._check()
            except Exception as e:
                _log.error(f"看门狗检查异常: {e}")

    def _check(self):
        now = time.time()
        with _WORKERS_LOCK:
            stale = []
            for serial, s in _WORKERS.items():
                # 只检查运行中的 worker
                if s.get("status") not in ("running", "connecting"):
                    continue
                hb = s.get("last_heartbeat", 0)
                if hb and now - hb > _HEARTBEAT_TIMEOUT:
                    stale.append(serial)
        for serial in stale:
            _log.error(f"[{serial}] 心跳超时 {now - get_worker_heartbeat(serial):.0f}s，标记卡死")
            _update_status(serial, status="error",
                           last_error=f"心跳超时 {_HEARTBEAT_TIMEOUT}s，worker 可能卡死")

    def stop(self):
        self._stop.set()


_watchdog = None


def start_watchdog():
    """启动看门狗（web_server 启动时调一次）。"""
    global _watchdog
    if _watchdog is None or not _watchdog.is_alive():
        _watchdog = _Watchdog()
        _watchdog.start()
        _log.info("心跳看门狗已启动")


def stop_watchdog():
    """停止看门狗。"""
    global _watchdog
    if _watchdog:
        _watchdog.stop()
        _watchdog = None


# ================== 通用 Worker 基类 ==================
class BaseWorker(threading.Thread):
    """通用 worker 基类。

    自动处理：STF 占用/释放、u2 连接、状态上报、异常捕获、stop 信号、心跳。
    子类只需实现 run_task(d) 方法，专注业务逻辑。

    生命周期（基类 run() 已封装，不要重写）：
      1. acquire 设备（STF 占用 + adb 连接）
      2. u2.connect 拿到 d
      3. 调用 setup(d)        ← 子类可选钩子，做初始化
      4. 调用子类 run_task(d)  ← 业务逻辑
      5. 调用 teardown(d)     ← 子类可选钩子，做清理
      6. finally: release 设备

    子类可选钩子：
      setup(d)          — run_task 前的初始化（如启动 app、检查登录）。默认空。
      teardown(d)       — run_task 后的清理（无论成功失败）。默认空。
      on_error(d, err)  — run_task 抛异常时回调。默认记录日志。

    子类可用工具方法：
      self.stopped()       — 判断是否被要求停止（循环里检查）
      self.set_action(s)   — 上报当前动作（同时刷新心跳）
      self.set_progress()  — 上报进度字段（同时刷新心跳）
      self.heartbeat()     — 手动刷新心跳（长操作中间调）

    进度上报规范（通用，适配任意 app）：
      self.set_progress(done=5, total=80, unit="视频")
        - done: 已完成数量（int）
        - total: 总数量（int）
        - unit: 计数单位（str，如 "视频"/"轮次"/"条"）
        - action_counts: 操作计数 dict（可选，如 {"like": 3, "comment": 1}）
        - 也可传任意业务字段，前端会原样展示在"进度"列
      示例：
        self.set_progress(done=5, total=80, unit="视频",
                          action_counts={"like": 3, "comment": 1})
      前端展示：进度条 + "5/80 视频" + 操作计数徽章

    运行时长终止（通用，适配任意 app）：
      self.max_duration — 最大运行时长（秒），0=不限时
      self.is_time_up() — 判断是否已达 max_duration
      子类在循环里检查 self.stopped() or self.is_time_up()
      到时自动停（自然结束，不是 error），状态=done

    异常分类（影响上层重试策略）：
      DeviceOfflineError — 设备掉线/不可用，不重试（换设备也没用）
      其他异常           — 可重试
    """

    def __init__(self, stf_client, serial, params=None, daemon=True):
        super().__init__(daemon=daemon)
        self.stf = stf_client
        self.serial = serial
        self.params = params or {}
        self._stop_flag = threading.Event()
        self.d = None  # u2.Device，run_task 里用
        # 通用进度字段（子类通过 set_progress 上报）
        self._progress = {"done": 0, "total": 0, "unit": "",
                          "action_counts": {}}
        # 通用运行时长终止（子类设 max_duration>0 启用）
        self.max_duration = 0           # 秒，0=不限时
        self._start_time = 0.0          # run_task 开始时间戳

    def _start_timer(self):
        """子类在 run_task 开头调用，启动运行时长计时。"""
        self._start_time = time.time()
        if self.max_duration > 0:
            # 预计结束时间（前端监控页展示），到点自动停止
            _update_status(self.serial, end_time=self._start_time + self.max_duration)

    def is_time_up(self):
        """是否已达最大运行时长。max_duration=0 时永远返回 False。"""
        if self.max_duration <= 0 or self._start_time <= 0:
            return False
        return (time.time() - self._start_time) >= self.max_duration

    def elapsed(self):
        """已运行时长（秒）。run_task 未启动时返回 0。"""
        if self._start_time <= 0:
            return 0
        return int(time.time() - self._start_time)

    def stop(self):
        """外部调用，请求停止（子类应在循环里检查 self.stopped()）。"""
        self._stop_flag.set()

    def stopped(self):
        """子类在循环里调用，判断是否被要求停止。"""
        return self._stop_flag.is_set()

    def set_action(self, action):
        """上报当前动作（子类用）。同时刷新心跳。"""
        _update_status(self.serial, current_action=action)

    def set_progress(self, **fields):
        """上报进度字段（子类用）。同时刷新心跳。

        通用字段（前端会解析展示）：
          done (int)         — 已完成数量
          total (int)        — 总数量
          unit (str)         — 计数单位（"视频"/"轮次"/"条"）
          action_counts (dict) — 操作计数 {"like": 3, "comment": 1}

        也可传任意业务字段，会合并到 progress 里，前端原样展示。

        示例：
          self.set_progress(done=5, total=80, unit="视频",
                            action_counts={"like": 3, "comment": 1})
        """
        self._progress.update(fields)
        # 同步到全局状态字典，前端通过 get_status 读取
        _update_status(self.serial, progress=dict(self._progress))

    def heartbeat(self):
        """手动刷新心跳。长操作中间调，避免看门狗误判卡死。"""
        _update_status(self.serial, last_heartbeat=time.time())

    # ================== 子类实现/可选钩子 ==================
    def run_task(self, d):
        """子类实现：拿到 u2.Device d，做具体业务。

        不需要管 STF 占用/释放、异常、状态上报，基类已处理。
        循环里记得检查 self.stopped() 以便响应停止请求。
        """
        raise NotImplementedError

    def setup(self, d):
        """run_task 前的初始化钩子。默认空，子类按需重写。

        典型用途：启动 app、检查登录状态、进入主页。
        抛异常会中断任务（走 on_error）。
        """

    def teardown(self, d):
        """run_task 后的清理钩子（无论成功失败都调）。默认空，子类按需重写。

        典型用途：退出 app、清理临时文件。
        抛异常只记日志，不影响主流程。
        """

    def on_error(self, d, error):
        """run_task 抛异常时的回调。默认空，子类按需重写。

        典型用途：截图保存现场、上报错误详情。
        """

    # ================== 基类主循环（不要重写） ==================
    def run(self):
        """基类主循环：不要重写。子类实现 run_task + 可选钩子。"""
        device = STFDevice(self.stf, serial=self.serial)
        _update_status(self.serial, status="connecting", stf_occupied=False,
                       last_error="", remote_adb_url="", model="",
                       current_action="")
        try:
            device.acquire()
            remote = device.remote_adb_url
            # u2.connect 可能 hang 住（atx-agent 无响应），用线程池+超时保护
            self.d = self._u2_connect_with_timeout(remote)
            if self.d is None:
                raise STFError(f"u2.connect 超时 {_U2_CONNECT_TIMEOUT}s，atx-agent 无响应", "offline")
            _log.info(f"[{self.serial}] u2 连接成功")
            model = ""
            try:
                info = self._u2_info_with_timeout()
                if info:
                    model = info.get("productName") or ""
            except Exception:
                pass
            _update_status(self.serial, status="running", stf_occupied=True,
                           remote_adb_url=remote, model=model)

            # setup 钩子
            try:
                self.setup(self.d)
            except Exception as e:
                _log.error(f"[{self.serial}] setup 异常: {e}")
                raise

            # 业务逻辑
            _log.info(f"[{self.serial}] 进入 run_task，开始执行业务逻辑")
            self.run_task(self.d)
            _log.info(f"[{self.serial}] run_task 执行完毕")

            # teardown 钩子（成功路径）
            try:
                self.teardown(self.d)
            except Exception as e:
                _log.warning(f"[{self.serial}] teardown 异常: {e}")

            if not self.stopped():
                _update_status(self.serial, status="done")

        except DeviceOfflineError as e:
            # 设备掉线，不重试
            _log.error(f"[{self.serial}] 设备离线: {e}")
            _update_status(self.serial, status="error", last_error=str(e))
            self._safe_on_error(e)
        except Exception as e:
            err_str = str(e)
            # 识别 Windows TCP 端口耗尽（WinError 10048）或 adb server 连接失败
            # 这类错误是临时性的，重试时需要更长退避等端口释放
            is_transient = ("10048" in err_str
                            or "connect to adb server failed" in err_str
                            or "Unable to connect to uiautomator2 server" in err_str)
            if is_transient:
                _log.error(f"[{self.serial}] ADB/u2 连接错误（可能端口耗尽）: {e}")
                _update_status(self.serial, status="error",
                               last_error=f"[transient] {e}")
            else:
                _log.error(f"[{self.serial}] 异常: {e}", exc_info=True)
                _update_status(self.serial, status="error", last_error=str(e))
            self._safe_on_error(e)
        finally:
            try:
                device.release()
            except Exception as e:
                _log.warning(f"[{self.serial}] release 异常: {e}")
            # 只在仍在运行/连接中时才标记 released；
            # done/error/failed 等终态保持不变，避免覆盖业务结果
            with _WORKERS_LOCK:
                cur = _WORKERS.get(self.serial, {}).get("status")
            if cur in ("running", "connecting"):
                _update_status(self.serial, stf_occupied=False, status="released")
            else:
                _update_status(self.serial, stf_occupied=False)

    def _u2_connect_with_timeout(self, remote):
        """u2.connect 带超时保护，避免 atx-agent 无响应时永久 hang。"""
        with ThreadPoolExecutor(max_workers=1) as pool:
            fut = pool.submit(u2.connect, remote)
            try:
                return fut.result(timeout=_U2_CONNECT_TIMEOUT)
            except FuturesTimeout:
                _log.error(f"[{self.serial}] u2.connect 超时 {_U2_CONNECT_TIMEOUT}s")
                return None
            except Exception as e:
                _log.error(f"[{self.serial}] u2.connect 异常: {e}")
                return None

    def _u2_info_with_timeout(self):
        """self.d.info 带超时保护，避免 atx-agent 无响应时永久 hang。"""
        with ThreadPoolExecutor(max_workers=1) as pool:
            fut = pool.submit(lambda: self.d.info)
            try:
                return fut.result(timeout=_U2_INFO_TIMEOUT)
            except FuturesTimeout:
                _log.warning(f"[{self.serial}] d.info 超时 {_U2_INFO_TIMEOUT}s，跳过型号获取")
                return None
            except Exception as e:
                _log.warning(f"[{self.serial}] d.info 异常: {e}")
                return None

    def _safe_on_error(self, error):
        """安全调用 on_error 钩子，不让钩子异常影响主流程。"""
        try:
            self.on_error(self.d, error)
        except Exception as e:
            _log.warning(f"[{self.serial}] on_error 钩子异常: {e}")