"""Worker 注册表：多 worker 池的注册、心跳与存活判断。

零第三方依赖，只用 Python 标准库。

设计：
- 注册表存在控制面数据库的 workers 表里（worker 与 server 共用
  同一个库，所以 worker 直接写库注册，不需要为此再发 HTTP；
  HTTP 侧的 GET /workers 与 POST /workers/heartbeat 只是同一张
  表的 API 面，给有 API key 的外部 worker 用）。
- worker_id 默认由调用方给（生产 worker 应配稳定的 --worker-id）；
  未给时由 worker.py 用“主机名-进程号”生成——重启后进程号会变，
  旧行会自然变 stale（alive=false），不要用随机 uuid 让注册表
  每次重启都多一行无法对应到同一台机器。
- 心跳本身就是注册：heartbeat 时 worker_id 不存在就 INSERT，
  存在就 UPDATE，绝不因“未注册”返回 404 报错。
- alive 定义：last_seen 在阈值内（默认 120 秒，可经环境变量
  CLOUD_WORKER_ALIVE_SECONDS 调，测试可传参覆盖）。
- 本模块不存、不打任何密钥：注册信息只有 worker_id、executor、
  capacity、last_seen，没有任何凭证字段。
"""

import os
import time

from cloud.db import connect as _db_connect

DEFAULT_ALIVE_SECONDS = 120


def alive_threshold_seconds(override=None):
    """存活阈值：显式参数 > 环境变量 > 默认 120 秒。"""
    if override is not None:
        try:
            return float(override)
        except (TypeError, ValueError):
            return float(DEFAULT_ALIVE_SECONDS)
    try:
        return float(os.environ.get(
            "CLOUD_WORKER_ALIVE_SECONDS") or DEFAULT_ALIVE_SECONDS)
    except (TypeError, ValueError):
        return float(DEFAULT_ALIVE_SECONDS)


def ensure_workers_table(db_path):
    """建 workers 表（不存在时建，已存在时不动，老库补精度）。

    注意：last_seen/created_at 必须是 DOUBLE PRECISION，不能用 REAL。
    SQLite 的 REAL 是 8 字节双精度，没问题；但 Postgres 的 REAL 是
    4 字节 float4，只能存约 7 位有效数字，epoch 秒（10 位）会被
    舍入到 128 秒一档，导致存活判断误差可达 ±64 秒、甚至出现负 age
    （2026-10-26 主机实测：last_seen 比 NOW() 还大 20 秒）。
    新表直接用 DOUBLE PRECISION；老库（已是 REAL 列）用 ALTER
    TYPE 升到 DOUBLE PRECISION（float4 -> float8 无损，已存的舍入值
    无法恢复，但之后的心跳不再舍入）。SQLite 下 ALTER TYPE 语法
    不支持，靠 try/except 跳过——SQLite 的 REAL 本来就是双精度，
    不需要升。
    """
    conn = _db_connect(str(db_path))
    try:
        conn.execute(
            "CREATE TABLE IF NOT EXISTS workers("
            "worker_id TEXT PRIMARY KEY, "
            "executor TEXT, "
            "capacity INTEGER, "
            "running INTEGER, "
            "last_seen DOUBLE PRECISION, "
            "created_at DOUBLE PRECISION)"
        )
        conn.commit()
        _ensure_running_column(conn)
        # 老库升精度只在确实是 Postgres REAL 列时做一次：每次心跳都
        # ALTER 会反复拿 ACCESS EXCLUSIVE 锁，没必要。先查
        # information_schema（SQLite 下这张表不存在，查询失败即跳过，
        # SQLite 的 REAL 本来就是双精度，不需要升）。
        for col in ("last_seen", "created_at"):
            try:
                row = conn.execute(
                    "SELECT data_type FROM information_schema.columns "
                    "WHERE table_name = 'workers' AND column_name = ?",
                    (col,)).fetchone()
            except Exception:
                try:
                    conn.rollback()
                except Exception:
                    pass
                continue
            if not row or row[0] != "real":
                continue
            try:
                conn.execute(
                    "ALTER TABLE workers ALTER COLUMN %s "
                    "TYPE DOUBLE PRECISION" % col)
                conn.commit()
            except Exception:
                try:
                    conn.rollback()
                except Exception:
                    pass
    finally:
        conn.close()


def _ensure_running_column(conn):
    """老库补 running 列（新表建表时已含该列）。

    running 是 worker 上次心跳上报的当前占用槽数（启动中占位
    也算，见 cloud/caps.py）；列缺失时 _worker_to_dict 默认 0，
    所以补列失败只回滚、不抛，绝不炸注册与心跳。
    SQLite 用 PRAGMA 查列，Postgres 用 information_schema。
    """
    try:
        rows = conn.execute("PRAGMA table_info(workers)").fetchall()
        names = set()
        for row in rows:
            try:
                names.add(row[1])
            except (IndexError, TypeError):
                pass
        if rows and "running" not in names:
            conn.execute(
                "ALTER TABLE workers ADD COLUMN running INTEGER DEFAULT 0")
            conn.commit()
        return
    except Exception:
        try:
            conn.rollback()
        except Exception:
            pass
    try:
        row = conn.execute(
            "SELECT COUNT(*) FROM information_schema.columns "
            "WHERE table_name = 'workers' AND column_name = 'running'"
            ).fetchone()
        exists = bool(row and row[0])
    except Exception:
        try:
            conn.rollback()
        except Exception:
            pass
        return
    if not exists:
        try:
            conn.execute(
                "ALTER TABLE workers ADD COLUMN IF NOT EXISTS "
                "running INTEGER DEFAULT 0")
            conn.commit()
        except Exception:
            try:
                conn.rollback()
            except Exception:
                pass


def register_worker(db_path, worker_id, executor=None, capacity=1,
                    running=None):
    """注册（或刷新）一个 worker：不存在 INSERT，存在 UPDATE。

    返回注册行 dict。worker_id 为空时抛 ValueError，调用方
    （server 的 heartbeat 端点）据此返回 400。
    """
    if not worker_id or not isinstance(worker_id, str):
        raise ValueError("worker_id required")
    ensure_workers_table(db_path)
    now = time.time()
    conn = _db_connect(str(db_path))
    try:
        cur = conn.execute(
            "UPDATE workers SET executor = ?, capacity = ?,"
            " running = ?, last_seen = ? WHERE worker_id = ?",
            (executor, capacity, running if running is not None else 0,
             now, worker_id),
        )
        if cur.rowcount == 0:
            conn.execute(
                "INSERT INTO workers(worker_id, executor, capacity,"
                " running, last_seen, created_at)"
                " VALUES(?, ?, ?, ?, ?, ?)",
                (worker_id, executor, capacity,
                 running if running is not None else 0, now, now),
            )
        conn.commit()
        row = conn.execute(
            "SELECT * FROM workers WHERE worker_id = ?",
            (worker_id,)).fetchone()
    finally:
        conn.close()
    return _worker_to_dict(row)


def heartbeat_worker(db_path, worker_id, executor=None, capacity=None,
                     running=None):
    """心跳：就是带 last_seen 刷新的注册。

    executor/capacity 未给时保留旧值（已注册的 worker 只刷时间，
    不把 executor 洗成 None）；未注册时按给定值（或默认）注册。
    """
    if not worker_id or not isinstance(worker_id, str):
        raise ValueError("worker_id required")
    ensure_workers_table(db_path)
    now = time.time()
    conn = _db_connect(str(db_path))
    try:
        row = conn.execute(
            "SELECT * FROM workers WHERE worker_id = ?",
            (worker_id,)).fetchone()
        if row is None:
            conn.execute(
                "INSERT INTO workers(worker_id, executor, capacity,"
                " running, last_seen, created_at)"
                " VALUES(?, ?, ?, ?, ?, ?)",
                (worker_id, executor, capacity if capacity is not None
                 else 1, running if running is not None else 0, now, now),
            )
        else:
            old = _worker_to_dict(row)
            conn.execute(
                "UPDATE workers SET executor = ?, capacity = ?,"
                " running = ?, last_seen = ? WHERE worker_id = ?",
                (executor if executor is not None else old["executor"],
                 capacity if capacity is not None else old["capacity"],
                 running if running is not None else old["running"],
                 now, worker_id),
            )
        conn.commit()
        row = conn.execute(
            "SELECT * FROM workers WHERE worker_id = ?",
            (worker_id,)).fetchone()
    finally:
        conn.close()
    return _worker_to_dict(row)


def list_workers(db_path, alive_seconds=None):
    """列出全部 worker，每行含 alive（last_seen 在阈值内）。"""
    threshold = alive_threshold_seconds(alive_seconds)
    ensure_workers_table(db_path)
    now = time.time()
    conn = _db_connect(str(db_path))
    try:
        rows = conn.execute(
            "SELECT * FROM workers ORDER BY worker_id").fetchall()
    finally:
        conn.close()
    out = []
    for row in rows:
        item = _worker_to_dict(row)
        last = item.get("last_seen")
        item["alive"] = bool(
            last is not None and (now - last) <= threshold)
        out.append(item)
    return out


def _worker_to_dict(row):
    def _get(name):
        try:
            return row[name]
        except (KeyError, IndexError, TypeError):
            return None
    running = _get("running")
    return {
        "worker_id": _get("worker_id"),
        "executor": _get("executor"),
        "capacity": _get("capacity"),
        "running": running if running is not None else 0,
        "last_seen": _get("last_seen"),
        "created_at": _get("created_at"),
    }