"""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，没有任何凭证字段。
- host/pid/address（本轨新增）：新库建表直接含三列；老库由
  ensure_workers_table 幂等补列，已有行新列为空、不编造值。
  两路界线：SQLite 文件库与 PG 生产库共用 ensure_workers_table
  这一份迁移代码（PG 下 connect() 返回适配连接，PRAGMA 经
  适配层翻译为 information_schema 查询）；PG 批量建库
  （POSTGRES_SCHEMA / init_postgres_schema）刻意不含 workers
  表——workers 是易失心跳数据，见 migrations.py 的
  UNMIGRATED_TABLES，把它加进批量 schema 会撞 P2 迁移对齐
  （MIGRATION_TABLES 必须等于 schema 表集合，本轨实测 3 项
  变红后已回退），不要再加。
  心跳刷新 host/pid 时未给保留旧值；address 没有新来源，只存
  调用方显式给的值，未给保留旧值。本模块绝不按 worker 名或
  前缀枚举/删除容器：失联清理只做判定（见
  classify_stale_workers），判定为 kill 的唯一可杀对象是该行
  的 pid 本人，且必须是本机 + pid 已死（pid_verified）。
  执行路径（本轨新增，见 reap_stale_workers）：复用
  classify_stale_workers 的判定，不另造判定函数；只对
  action == "kill" 的行调 kill_fn，且执行时先用同一 checker
  再验一次已死（防判定快照与执行之间 pid 被复用），杀前
  写审计（审计写失败则不杀，fail-closed），杀后用同一
  checker 复核 pid 已死；异机、无 pid、pid 存活或未知一律
  只列不杀。执行路径不枚举容器、不按名/前缀匹配、不碰
  基础设施容器、不删 workers 注册行：杀的是 pid 本人，
  注册行的去留是另一条决策（删行会改 /workers 列表语义），
  本轨不擅自删。亦不新增 HTTP 路由（会改 v1.1 契约）：调
  用方（管理员/清扫器）在有权限的上下文里直接调本函数。
"""

import os
import signal
import socket
import time

from cloud.db import connect as _db_connect

DEFAULT_ALIVE_SECONDS = 120

# workers 表新增列：(列名, 类型)。pid 用 BIGINT：SQLite 下为
# INTEGER 亲和性，Postgres 下为 8 字节整数，均可容纳进程号；
# host/address 为 TEXT，不存任何凭证。
WORKER_HOSTPID_COLUMNS = (
    ("host", "TEXT"),
    ("pid", "BIGINT"),
    ("address", "TEXT"),
)


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, "
            "host TEXT, "
            "pid BIGINT, "
            "address TEXT, "
            "last_seen DOUBLE PRECISION, "
            "created_at DOUBLE PRECISION)"
        )
        conn.commit()
        _ensure_running_column(conn)
        _ensure_hostpid_columns(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 _ensure_hostpid_columns(conn):
    """老库补 host/pid/address 列（新表建表时已含三列）。

    幂等保证：先查现有列名，只补缺失列，多次跑是 no-op；已有
    行新列保持 NULL，不编造 host/pid/address。SQLite 与 PG
    共用本函数：SQLite 走 PRAGMA table_info，Postgres 走适配
    层翻译后的同款 PRAGMA（见 cloud/db.py 的 PRAGMA 翻译，调用
    方取 row[1] 即列名，与 init_postgres_schema 补列同款）；
    PRAGMA 查不到列时再走 information_schema 逐列 COUNT 兜底
    （Postgres 语法 ADD COLUMN IF NOT EXISTS，SQLite 下
    PRAGMA 本就可用，不会走到这条兜底）。补列失败只回滚、
    不抛，绝不炸注册与心跳。PG 批量建库路径不适用本函数：
    POSTGRES_SCHEMA 刻意不含 workers（易失心跳数据，见模块
    头注释），PG 生产库的 workers 表同样由本函数懒建/补列。
    """
    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:
            for col, col_type in WORKER_HOSTPID_COLUMNS:
                if col not in names:
                    conn.execute(
                        "ALTER TABLE workers ADD COLUMN %s %s"
                        % (col, col_type))
            conn.commit()
            return
    except Exception:
        try:
            conn.rollback()
        except Exception:
            pass
    for col, col_type in WORKER_HOSTPID_COLUMNS:
        try:
            row = conn.execute(
                "SELECT COUNT(*) FROM information_schema.columns "
                "WHERE table_name = 'workers' AND column_name = '%s'"
                % col).fetchone()
            exists = bool(row and row[0])
        except Exception:
            try:
                conn.rollback()
            except Exception:
                pass
            continue
        if not exists:
            try:
                conn.execute(
                    "ALTER TABLE workers ADD COLUMN IF NOT EXISTS "
                    "%s %s" % (col, col_type))
                conn.commit()
            except Exception:
                try:
                    conn.rollback()
                except Exception:
                    pass


def _normalize_pid(pid):
    """归一进程号：None -> None；bool 或不可转 int -> ValueError。

    pid 列存整数，不存字符串形式的进程号；调用方传错类型时
    明确报错，不静默写成 0 或把旧值洗掉。
    """
    if pid is None:
        return None
    if isinstance(pid, bool):
        raise ValueError("pid must be an integer")
    try:
        value = int(pid)
    except (TypeError, ValueError):
        raise ValueError("pid must be an integer")
    if value <= 0:
        raise ValueError("pid must be a positive integer")
    return value


def pid_is_dead(pid):
    """本机进程是否明确已死：os.kill(pid, 0) 的语义。

    ProcessLookupError -> 已死（True）；PermissionError -> 进程
    存在但无权发信号，不能证明已死（False）；其它 OSError /
    非法 pid -> 不能证明已死（False）。绝不把“查不到/无权”
    当成已死，这是 pid_verified 不误杀的底线。
    """
    try:
        value = _normalize_pid(pid)
    except ValueError:
        return False
    if value is None:
        return False
    try:
        os.kill(value, 0)
    except ProcessLookupError:
        return True
    except PermissionError:
        return False
    except OSError:
        return False
    return False


def classify_stale_workers(db_path, stale_seconds=None, local_host=None,
                           pid_is_dead_fn=None, now=None):
    """失联 worker 判定（只判定、不杀任何进程或容器）。

    返回全部 worker 行，每行含 action 与 reason：
    - action == "kill" 当且仅当同时满足：age 严格大于阈值
      （恰等于阈值不杀）且 pid_verified（行 host 等于本机 host、
      pid 为已知正整数、且 pid_is_dead_fn(pid) 为 True）。
    - 其余一律 action == "list_only"：未失联、恰等于阈值、
      host 不是本机、pid 缺失、pid 还活着或无法证明已死，
      都只列不杀。
    - 本函数不枚举容器、不按 worker_id 前缀匹配、不碰任何
      基础设施容器：kill 项的唯一可杀对象就是该行的 pid 本人，
      调用方若要清理也必须按 pid 精确杀，绝不按名/前缀删。
    - stale_seconds 未给时复用 alive 阈值语义（显式参数 >
      环境变量 > 120 秒）；local_host 未给时为本机主机名；
      pid_is_dead_fn 未给时为本机 pid_is_dead（测试可注入）。
    """
    threshold = alive_threshold_seconds(stale_seconds)
    host = local_host if local_host is not None else socket.gethostname()
    checker = pid_is_dead_fn if pid_is_dead_fn is not None else pid_is_dead
    moment = time.time() if now is None else float(now)
    ensure_workers_table(db_path)
    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")
        age = (moment - last) if last is not None else None
        stale = bool(age is not None and age > threshold)
        verified = False
        if stale:
            row_pid = item.get("pid")
            if (item.get("host") == host and row_pid is not None
                    and checker(row_pid)):
                verified = True
        if verified:
            action, reason = "kill", "pid_verified_dead"
        elif age is None:
            action, reason = "list_only", "no_last_seen"
        elif age == threshold:
            action, reason = "list_only", "age_equals_threshold"
        elif age <= threshold:
            action, reason = "list_only", "not_stale"
        elif item.get("host") != host:
            action, reason = "list_only", "host_mismatch"
        elif item.get("pid") is None:
            action, reason = "list_only", "pid_missing"
        else:
            action, reason = "list_only", "pid_alive_or_unknown"
        item["age_seconds"] = age
        item["stale"] = stale
        item["pid_verified"] = verified
        item["action"] = action
        item["reason"] = reason
        out.append(item)
    return out


def _default_kill(pid):
    """默认真杀：只给该 pid 本人发 SIGKILL，不发进程组。

    pid 已死时 os.kill 抛 ProcessLookupError，调用方
    （reap_stale_workers）把它当“已死、无需再杀”处理，
    不是失败。绝不把信号发给进程组（会误伤同组无关进程），
    绝不按名/前缀找进程。
    """
    os.kill(pid, signal.SIGKILL)


def _default_audit(db_path, action, detail):
    """默认审计写 audit_log：user_id=system、project_id=""。

    表不存在时先建（PG 与 SQLite 各一套 DDL，同
    cloud/automations.py 的兜底写法：PG DDL 在 SQLite 下
    失败即回滚再试 SQLite DDL，PG 适配层下第一条即成功）。
    写失败抛给调用方：reap 的杀前审计失败时不杀
    （fail-closed），杀后审计失败时只记 audit_after=False。
    """
    conn = _db_connect(str(db_path))
    try:
        for ddl in (
                "CREATE TABLE IF NOT EXISTS audit_log("
                "id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, "
                "user_id TEXT, project_id TEXT, action TEXT, detail TEXT, "
                "created_at DOUBLE PRECISION)",
                "CREATE TABLE IF NOT EXISTS audit_log("
                "id INTEGER PRIMARY KEY AUTOINCREMENT, "
                "user_id TEXT, project_id TEXT, action TEXT, detail TEXT, "
                "created_at REAL)"):
            try:
                conn.execute(ddl)
                conn.execute(
                    "INSERT INTO audit_log(user_id, project_id, action,"
                    " detail, created_at) VALUES(?, ?, ?, ?, ?)",
                    ("system", "", action, detail, time.time()))
                conn.commit()
                return
            except Exception:
                try:
                    conn.rollback()
                except Exception:
                    pass
        raise RuntimeError("audit_log write failed")
    finally:
        conn.close()


def reap_stale_workers(db_path, stale_seconds=None, local_host=None,
                       pid_is_dead_fn=None, kill_fn=None, audit_fn=None,
                       now=None):
    """失联 worker 执行路径：复用判定，只杀 pid_verified。

    流程（每行）：
    - 先调 classify_stale_workers（不另造判定）：action 不是
      kill 的行一律只列不杀，executed=False、outcome="listed"，
      绝不调 kill_fn、不写杀审计。
    - kill 行执行时先用同一 checker 再验一次 pid 已死：若此时
      已不死（判定后 pid 被复用等），跳过不杀，
      outcome="skipped_pid_alive_at_execution"，不调 kill_fn。
    - 杀前写审计：audit_fn 未给时走 _default_audit 写
      audit_log（action=worker_pid_kill）；审计写失败则不杀，
      outcome="audit_failed"（fail-closed）。
    - 杀：kill_fn 未给时走 _default_kill（SIGKILL 只给该 pid
      本人）；ProcessLookupError 当已死处理，不是失败；其它
      OSError 记入 kill_error，仍以后续复核定成败。
    - 杀后用同一 checker 复核：已死 -> outcome="killed"；
      还活着 -> outcome="kill_failed_pid_alive"；随后写
      杀后审计（action=worker_pid_kill_verified，detail 含
      outcome），审计成败只记 audit_after，不改 outcome。
    - 本函数不枚举容器、不按名/前缀匹配、不碰基础设施容器、
      不删 workers 行；返回每行的判定字段 + executed /
      outcome / verified_dead / audit_before / audit_after /
      kill_error。verified_dead 仅在杀后复核后才有真值，
      未执行时为 None（不是“还活着”的断言）。
    """
    rows = classify_stale_workers(
        db_path, stale_seconds=stale_seconds, local_host=local_host,
        pid_is_dead_fn=pid_is_dead_fn, now=now)
    checker = pid_is_dead_fn if pid_is_dead_fn is not None else pid_is_dead
    killer = kill_fn if kill_fn is not None else _default_kill
    results = []
    for row in rows:
        result = dict(row)
        result["executed"] = False
        result["verified_dead"] = None
        result["audit_before"] = False
        result["audit_after"] = False
        result["kill_error"] = None
        if row.get("action") != "kill":
            result["outcome"] = "listed"
            results.append(result)
            continue
        pid = row.get("pid")
        if pid is None or not checker(pid):
            result["outcome"] = "skipped_pid_alive_at_execution"
            result["verified_dead"] = False
            results.append(result)
            continue
        detail = ("worker_id=%s host=%s pid=%s reason=%s"
                  % (row.get("worker_id"), row.get("host"), pid,
                     row.get("reason")))
        try:
            if audit_fn is not None:
                audit_fn("worker_pid_kill", detail)
            else:
                _default_audit(db_path, "worker_pid_kill", detail)
            result["audit_before"] = True
        except Exception:
            result["outcome"] = "audit_failed"
            results.append(result)
            continue
        result["executed"] = True
        try:
            killer(pid)
        except ProcessLookupError:
            pass
        except Exception as exc:
            result["kill_error"] = str(exc)
        verified = bool(checker(pid))
        result["verified_dead"] = verified
        result["outcome"] = (
            "killed" if verified else "kill_failed_pid_alive")
        after_detail = detail + " outcome=%s verified_dead=%s" % (
            result["outcome"], verified)
        try:
            if audit_fn is not None:
                audit_fn("worker_pid_kill_verified", after_detail)
            else:
                _default_audit(
                    db_path, "worker_pid_kill_verified", after_detail)
            result["audit_after"] = True
        except Exception:
            result["audit_after"] = False
        results.append(result)
    return results


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

    返回注册行 dict。worker_id 为空时抛 ValueError，调用方
    （server 的 heartbeat 端点）据此返回 400。新列语义：新行
    按给定值写（未给为 NULL）；已存在行上给定值覆盖、未给
    保留旧值——绝不因一次不带 host/pid 的注册把已有的身份
    洗成 NULL，也不编造 host/pid/address。
    """
    if not worker_id or not isinstance(worker_id, str):
        raise ValueError("worker_id required")
    pid = _normalize_pid(pid)
    ensure_workers_table(db_path)
    now = time.time()
    conn = _db_connect(str(db_path))
    try:
        old_row = conn.execute(
            "SELECT * FROM workers WHERE worker_id = ?",
            (worker_id,)).fetchone()
        if old_row is not None:
            old = _worker_to_dict(old_row)
            conn.execute(
                "UPDATE workers SET executor = ?, capacity = ?,"
                " running = ?, host = ?, pid = ?, address = ?,"
                " last_seen = ? WHERE worker_id = ?",
                (executor, capacity,
                 running if running is not None else 0,
                 host if host is not None else old["host"],
                 pid if pid is not None else old["pid"],
                 address if address is not None else old["address"],
                 now, worker_id),
            )
        else:
            conn.execute(
                "INSERT INTO workers(worker_id, executor, capacity,"
                " running, host, pid, address, last_seen, created_at)"
                " VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (worker_id, executor, capacity,
                 running if running is not None else 0,
                 host, pid, address, 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, host=None, pid=None, address=None):
    """心跳：就是带 last_seen 刷新的注册。

    executor/capacity 未给时保留旧值（已注册的 worker 只刷时间，
    不把 executor 洗成 None）；未注册时按给定值（或默认）注册。
    host/pid/address 同款保留语义：给定则刷新（本机 worker 心跳
    应传本机 host 与本进程 pid，见 worker._registry_heartbeat），
    未给则保留旧值；address 没有新来源，绝不由本函数推导
    （例如不从 worker_id 或 socket 猜地址），只存显式给的值。
    """
    if not worker_id or not isinstance(worker_id, str):
        raise ValueError("worker_id required")
    pid = _normalize_pid(pid)
    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, host, pid, address, last_seen, created_at)"
                " VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (worker_id, executor, capacity if capacity is not None
                 else 1, running if running is not None else 0,
                 host, pid, address, now, now),
            )
        else:
            old = _worker_to_dict(row)
            conn.execute(
                "UPDATE workers SET executor = ?, capacity = ?,"
                " running = ?, host = ?, pid = ?, address = ?,"
                " 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"],
                 host if host is not None else old["host"],
                 pid if pid is not None else old["pid"],
                 address if address is not None else old["address"],
                 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,
        "host": _get("host"),
        "pid": _get("pid"),
        "address": _get("address"),
        "last_seen": _get("last_seen"),
        "created_at": _get("created_at"),
    }