"""Worker 健康池 + 空闲派发。

- 后台 asyncio 任务周期性探测每个 worker 的 /health
- 连续失败 N 次 → 下线；连续成功 N 次 → 上线
- 每个 worker 一个 busy 标志（per_worker_concurrency=1）
- acquire_worker() 从在线池挑空闲 worker；全忙排队；池空→NoWorkerAvailable
- release_worker() 释放 worker
- 两级排队（全局槽 / worker 池）均支持优先级：priority 大者插到队列最前，
  同优先级 FIFO；已有等待者时新请求不得越过（严格按优先级排序）
"""

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import httpx

logger = logging.getLogger("gateway.pool")


def _cooldown_seconds() -> float:
    """被动摘除后的隔离期时长（秒），配置 health_check.cooldown_seconds，默认 60。"""
    try:
        from gateway.config import get_config
        return float(get_config()["health_check"].get("cooldown_seconds", 60))
    except Exception:  # noqa: BLE001
        return 60.0


# ---------------------------------------------------------------------------
# 异常
# ---------------------------------------------------------------------------

class NoWorkerAvailable(Exception):
    """无可用 worker（池空或全忙排队超时）。"""
    pass


# ---------------------------------------------------------------------------
# 数据结构
# ---------------------------------------------------------------------------

@dataclass
class WorkerState:
    """单个 worker 的运行时状态。"""
    url: str
    online: bool = False  # 初始 offline，等健康检查通过后上线
    busy: bool = False
    primary: bool = True  # 主 worker（True 优先派发）；全部主 worker 下线才用备组
    consecutive_failures: int = 0
    consecutive_successes: int = 0
    offline_until: float = 0.0  # 被动摘除后的隔离截止时刻（time.monotonic() 基准）


# ---------------------------------------------------------------------------
# 全局状态
# ---------------------------------------------------------------------------

_workers: Dict[str, WorkerState] = {}
_pool_condition: Optional[asyncio.Condition] = None
_health_task: Optional[asyncio.Task] = None
_shutdown_event: Optional[asyncio.Event] = None

# worker 池等待者：[(-priority, 入队序号, future)]，按 (-priority, 序号) 升序排列。
# 等待者在持有 _pool_condition 锁时增删；release_worker() 唤醒队首（优先级最高者）。
_worker_waiters: List[tuple] = []
_waiter_seq: int = 0

# 全局串行槽：同一时间最多 max_global_concurrency 个请求进入 worker 派发，其余排队。
# ⚠️ 依赖单进程 uvicorn 部署（hair-gateway.service 无 --workers）。
#    若改多 worker / gunicorn，进程内信号量会失效，需换成跨进程锁（文件锁/Redis）。
_global_sem: Optional["_PrioritySemaphore"] = None
_global_max: int = 1
_global_busy: int = 0
_global_waiting: int = 0


# ---------------------------------------------------------------------------
# 优先级信号量
# ---------------------------------------------------------------------------

class _PrioritySemaphore:
    """带插队优先级的信号量（替代 asyncio.Semaphore，后者不支持优先级）。

    等待者按 (-priority, 入队序号) 排序：priority 大者先获得槽位（插队到最前），
    同优先级 FIFO。仅当无等待者时才允许直接获取——防止新请求越过排队者
    （asyncio.Semaphore 本身允许这种插队，这里收紧以保证优先级语义）。
    """

    def __init__(self, value: int) -> None:
        self._value = value
        self._waiters: List[tuple] = []  # [(-priority, seq, future)]
        self._seq = 0

    async def acquire(self, priority: int = 0) -> None:
        if self._value > 0 and not self._waiters:
            self._value -= 1
            return
        loop = asyncio.get_running_loop()
        fut = loop.create_future()
        self._seq += 1
        entry = (-priority, self._seq, fut)
        idx = 0
        while idx < len(self._waiters) and self._waiters[idx][:2] <= entry[:2]:
            idx += 1
        self._waiters.insert(idx, entry)
        if priority > 0 and idx < len(self._waiters) - 1:
            logger.info("优先级请求插队到全局队列第 %d 位（超越 %d 个等待者）",
                        idx + 1, len(self._waiters) - 1 - idx)
        try:
            await fut
        except asyncio.CancelledError:
            if fut.done() and not fut.cancelled():
                # 已被 release 授予槽位、但本任务随后被取消 → 归还槽位
                self.release()
            else:
                try:
                    self._waiters.remove(entry)
                except ValueError:
                    pass
            raise

    def release(self) -> None:
        while self._waiters:
            _, _, fut = self._waiters.pop(0)
            if not fut.done():
                fut.set_result(None)
                return
        self._value += 1


# ---------------------------------------------------------------------------
# 健康检查后台任务
# ---------------------------------------------------------------------------

def _try_revive(w: WorkerState, healthy_threshold: int) -> bool:
    """探测成功后是否允许将该 worker 复活上线。

    条件：未在线 + 连续成功达到阈值 + **不在被动摘除隔离期内**
    （offline_until 之前即使 /health 探测成功也压着不上线，防横跳）。
    """
    if w.online or w.consecutive_successes < healthy_threshold:
        return False
    if time.monotonic() < w.offline_until:
        logger.debug("Worker %s 探测成功但处于隔离期，暂不复活", w.url)
        return False
    return True


async def _health_check_loop(cfg: dict) -> None:
    """后台循环：周期性探测所有 worker 健康状态。"""
    hc_cfg = cfg["health_check"]
    interval = hc_cfg["interval_seconds"]
    unhealthy_threshold = hc_cfg["unhealthy_threshold"]
    healthy_threshold = hc_cfg["healthy_threshold"]
    token = cfg["shared_password"]

    logger.info(
        "健康检查循环启动 | 间隔=%ds | 下线阈值=%d | 上线阈值=%d | 隔离期=%.0fs | workers=%d",
        interval, unhealthy_threshold, healthy_threshold, _cooldown_seconds(), len(_workers),
    )

    async with httpx.AsyncClient() as client:
        while not _shutdown_event.is_set():
            for w in _workers.values():
                url = f"{w.url}{hc_cfg['path']}"
                try:
                    resp = await client.get(
                        url,
                        headers={"X-Internal-Token": token},
                        timeout=hc_cfg["timeout_seconds"],
                    )
                    if resp.status_code == 200:
                        w.consecutive_failures = 0
                        w.consecutive_successes += 1
                        if _try_revive(w, healthy_threshold):
                            w.online = True
                            logger.info("✅ Worker 上线: %s", w.url)
                    else:
                        w.consecutive_successes = 0
                        w.consecutive_failures += 1
                        if w.consecutive_failures >= unhealthy_threshold and w.online:
                            w.online = False
                            logger.warning("⚠ Worker 下线: %s（HTTP %d，连续失败 %d 次）",
                                           w.url, resp.status_code, w.consecutive_failures)
                except Exception as exc:
                    w.consecutive_successes = 0
                    w.consecutive_failures += 1
                    if w.consecutive_failures >= unhealthy_threshold and w.online:
                        w.online = False
                        logger.warning("⚠ Worker 下线: %s（%s，连续失败 %d 次）",
                                       w.url, exc, w.consecutive_failures)

            # 等待下一次探测（支持快速关闭）
            try:
                await asyncio.wait_for(_shutdown_event.wait(), timeout=interval)
                break  # shutdown signaled
            except asyncio.TimeoutError:
                pass  # 正常的 interval 到期

    logger.info("健康检查循环已停止")


# ---------------------------------------------------------------------------
# 初始化 / 关闭
# ---------------------------------------------------------------------------

async def init_pool(cfg: dict) -> None:
    """初始化 worker 池并启动健康检查后台任务。"""
    global _workers, _pool_condition, _health_task, _shutdown_event

    # 主 worker 分组：配置了 primary_workers 则按其标记，否则全部视为主 worker
    primary_set = set(cfg.get("primary_workers") or cfg["workers"])

    _workers = {url: WorkerState(url=url, primary=url in primary_set)
                for url in cfg["workers"]}
    n_primary = sum(1 for w in _workers.values() if w.primary)
    _pool_condition = asyncio.Condition()
    _shutdown_event = asyncio.Event()

    # 全局串行槽（独立于 worker busy 标志）
    init_global_slot(cfg)

    # 立即做一轮健康检查以快速上线
    hc_cfg = cfg["health_check"]
    token = cfg["shared_password"]
    async with httpx.AsyncClient() as client:
        for w in _workers.values():
            try:
                resp = await client.get(
                    f"{w.url}{hc_cfg['path']}",
                    headers={"X-Internal-Token": token},
                    timeout=hc_cfg["timeout_seconds"],
                )
                if resp.status_code == 200:
                    w.consecutive_successes = 1
                    if hc_cfg["healthy_threshold"] <= 1:
                        w.online = True
                else:
                    w.consecutive_failures = 1
            except Exception:
                w.consecutive_failures = 1

    # 对于 healthy_threshold == 1 的情况，已在上面的检查中上线
    # 标记初始在线状态
    online_count = sum(1 for w in _workers.values() if w.online)
    logger.info("Worker 池初始化完成 | 总数=%d | 在线=%d | 主worker=%d",
                len(_workers), online_count, n_primary)

    # 启动后台健康检查
    _health_task = asyncio.create_task(_health_check_loop(cfg))


async def shutdown_pool() -> None:
    """关闭健康检查任务，释放资源。"""
    global _shutdown_event, _health_task
    if _shutdown_event:
        _shutdown_event.set()
    if _health_task:
        _health_task.cancel()
        try:
            await _health_task
        except asyncio.CancelledError:
            pass
        _health_task = None
    logger.info("Worker 池已关闭")


# ---------------------------------------------------------------------------
# 全局串行槽
# ---------------------------------------------------------------------------

def init_global_slot(cfg: dict) -> None:
    """初始化全局串行信号量。

    max_global_concurrency=N → 同一时间最多 N 个请求进入 worker 派发，其余排队。
    priority>0 的请求可插队到等待队列最前（见 _PrioritySemaphore）。
    依赖单进程 uvicorn 部署；多进程需换跨进程锁。
    """
    global _global_sem, _global_max, _global_busy, _global_waiting
    n = int(cfg.get("dispatch", {}).get("max_global_concurrency", 1))
    if n < 1:
        n = 1
    _global_sem = _PrioritySemaphore(n)
    _global_max = n
    _global_busy = 0
    _global_waiting = 0
    logger.info("全局并发槽初始化: max=%d（单进程生效）", n)


async def acquire_global_slot(cfg: dict, priority: int = 0) -> bool:
    """获取一个全局处理槽。True=获得；False=排队超时。

    priority>0 的请求插队到等待队列最前，同优先级 FIFO。
    超时由 dispatch.max_queue_wait_seconds 控制（需 < nginx proxy_read_timeout 600s）。
    """
    global _global_busy, _global_waiting
    if _global_sem is None:
        return True  # 未初始化，不限流
    timeout = float(cfg.get("dispatch", {}).get("max_queue_wait_seconds", 590))
    _global_waiting += 1
    try:
        await asyncio.wait_for(_global_sem.acquire(priority), timeout=timeout)
        _global_busy += 1
        return True
    except asyncio.TimeoutError:
        logger.warning("全局排队超时（%.1fs），拒绝请求 | waiting=%d", timeout, _global_waiting)
        return False
    finally:
        _global_waiting -= 1


def release_global_slot() -> None:
    """释放全局处理槽。"""
    global _global_busy
    if _global_sem is not None:
        _global_busy -= 1
        _global_sem.release()


# ---------------------------------------------------------------------------
# 派发
# ---------------------------------------------------------------------------

async def mark_worker_unhealthy(w: WorkerState) -> None:
    """被动标记：转发失败时立即将该 worker 下线，并进入隔离期。

    隔离期内即使 /health 探测成功也不复活（由 _try_revive 守卫），
    避免"摘除→8 秒复活→再失败"的横跳。
    """
    async with _pool_condition:
        if w.online:
            w.online = False
            w.consecutive_failures += 1
            cooldown = _cooldown_seconds()
            w.offline_until = time.monotonic() + cooldown
            logger.warning("⚠ Worker 被动下线: %s（转发失败，隔离 %.0fs）",
                           w.url, cooldown)


def _insert_worker_waiter(key: tuple, fut: asyncio.Future) -> tuple:
    """按 (-priority, 序号) 升序插入等待者（须在持有 _pool_condition 锁时调用）。"""
    entry = (key[0], key[1], fut)
    idx = 0
    while idx < len(_worker_waiters) and _worker_waiters[idx][:2] <= entry[:2]:
        idx += 1
    _worker_waiters.insert(idx, entry)
    return entry


def _wake_best_worker_waiter() -> None:
    """唤醒优先级最高的等待者（须在持有 _pool_condition 锁时调用）。"""
    while _worker_waiters:
        _, _, fut = _worker_waiters.pop(0)
        if not fut.done():
            fut.set_result(None)
            return


async def acquire_worker(cfg: dict, priority: int = 0) -> WorkerState:
    """从在线池中获取一个空闲 worker。

    主备分组：只要还有主 worker（primary）在线，就只在主组内挑空闲
    （全忙则排队等待，不用备组）；只有主组全部下线时才用备用 worker。
    全忙则排队等待，最多 queue_wait_seconds；池空立即抛异常。
    排队按 (-priority, 入队序号) 排序：priority 大者插队到最前，同优先级 FIFO；
    已有等待者时新请求不得越过（保证优先级语义，防饿死）。
    """
    queue_wait = cfg["dispatch"]["queue_wait_seconds"]
    deadline = time.monotonic() + queue_wait
    my_key: Optional[tuple] = None  # (-priority, seq)，被唤醒后重排仍用原序号
    woken = False  # 是否刚被 release_worker 唤醒（被唤醒者直接参与挑选）

    while True:
        async with _pool_condition:
            online = [w for w in _workers.values() if w.online]
            primary_online = [w for w in online if w.primary]
            # 主组有在线 → 只在主组内挑；主组全挂 → 才轮到备组
            candidates = primary_online if primary_online else online
            idle = [w for w in candidates if not w.busy]

            if idle and (woken or not _worker_waiters):
                # 取第一个空闲 worker
                w = idle[0]
                w.busy = True
                if not w.primary:
                    logger.info("主 worker 全部下线，使用备用 worker: %s", w.url)
                logger.debug("Worker 分配: %s", w.url)
                return w

            if not online:
                raise NoWorkerAvailable("后端服务暂不可用，请稍后重试")

            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise NoWorkerAvailable("后端服务繁忙，请稍后重试")

            if my_key is None:
                global _waiter_seq
                _waiter_seq += 1
                my_key = (-priority, _waiter_seq)
                first_enqueue = True
            else:
                first_enqueue = False
            fut = asyncio.get_running_loop().create_future()
            entry = _insert_worker_waiter(my_key, fut)
            if first_enqueue:
                pos = _worker_waiters.index(entry) + 1
                if not idle:
                    logger.info("所有可用 worker 全忙（%d 在线），等待 %.1fs...",
                                len(candidates), remaining)
                if priority > 0 and pos < len(_worker_waiters):
                    logger.info("优先级请求插队到 worker 队列第 %d 位（超越 %d 个等待者）",
                                pos, len(_worker_waiters) - pos)

        try:
            await asyncio.wait_for(fut, timeout=remaining)
            woken = True
        except asyncio.TimeoutError:
            async with _pool_condition:
                try:
                    _worker_waiters.remove(entry)
                except ValueError:
                    # 授予与取消竞态：已被弹出授予（worker 空闲），转手唤醒下一位
                    if fut.done() and not fut.cancelled():
                        _wake_best_worker_waiter()
            raise NoWorkerAvailable("后端服务繁忙，请稍后重试")


async def release_worker(w: WorkerState) -> None:
    """释放 worker，标记为空闲并唤醒优先级最高的等待者。"""
    async with _pool_condition:
        w.busy = False
        logger.debug("Worker 释放: %s", w.url)
        _wake_best_worker_waiter()


# ---------------------------------------------------------------------------
# 状态查询
# ---------------------------------------------------------------------------

def get_pool_status() -> dict:
    """返回当前池状态（供 /gateway-health 使用）。"""
    global_info = {
        "global_max": _global_max,
        "global_busy": _global_busy,
        "global_waiting": _global_waiting,
    }
    if not _workers:
        return {"total": 0, "healthy": 0, "busy": 0,
                "primary_total": 0, "primary_healthy": 0, **global_info}
    total = len(_workers)
    healthy = sum(1 for w in _workers.values() if w.online)
    busy = sum(1 for w in _workers.values() if w.busy)
    primary_total = sum(1 for w in _workers.values() if w.primary)
    primary_healthy = sum(1 for w in _workers.values() if w.primary and w.online)
    return {"total": total, "healthy": healthy, "busy": busy,
            "primary_total": primary_total, "primary_healthy": primary_healthy,
            **global_info}