#!/usr/bin/env python3
"""chv-daemon：客机运行时 daemon（阶段 2 三分之一）。

职责（只做这三件，不碰任务执行与同步）：
  1. 生命周期/心跳：每 10 秒写 /run/chv/heartbeat.json（atomic），
     含 pid、启动时间、seq、uptime、loadavg、内存、workspace
     余量、execd 任务数、sync 最近状态摘要。seq 跨重启连续
     （启动时读旧文件续号），宿主侧据此判断“心跳丢了还是
     daemon 重启了”。
  2. 资源上报：同上，心跳文件即上报载体（宿主经 SSH 读文件，
     不另开通道；控制通道仍是 SSH，见 executor.py 定稿）。
  3. 快照配合：Unix socket /run/chv/daemon.sock，op
     prepare-snapshot——问 execd 是否有运行中任务；有则
     拒绝（{ok:false, reason:task-running}），无则写
     /run/chv/snapshot-ready.json 标记并返回 ok；execd
     联系不上时同样拒绝（fail-closed，不写标记）。
     宿主侧快照流程暂未调它（阶段 2 不改 executor 快照流），
     接口先行、供宿主侧后续接线与本阶段 live 演练。

与既有分工的映射（不要第二套 agent）：
  - host.py 仍是唯一的任务/ codex 执行引擎（CLI，经 SSH 或
    execd 的 codex op 调）；daemon 从不跑 host.py。
  - chv-broker/chv-browser 仍是单元内外 IPC 与浏览器子代理，
    daemon 不碰它们的 socket。
  - sync 仍是 chv-sync.sh + chv-syncd 常驻循环，daemon 只读
    sync-status.json 作摘要，不调 git。

只用标准库。以 root 运行（读 /proc 与各服务状态文件无权限
问题；socket 0600 root:root，只接受本机 root 客户）。
"""
import json
import os
import signal
import socket
import sys
import threading
import time

HEARTBEAT_PATH = "/run/chv/heartbeat.json"
DAEMON_SOCK = "/run/chv/daemon.sock"
SNAPSHOT_READY_PATH = "/run/chv/snapshot-ready.json"
EXECD_SOCK = "/run/chv/execd.sock"
EXECD_TASKS = "/run/chv/execd-tasks.json"
SYNC_STATUS = "/run/chv/sync-status.json"
WORKSPACE = "/home/chv/workspace"
HEARTBEAT_INTERVAL = 10.0
MAX_REQUEST_BYTES = 65536


def _atomic_write(path, data: bytes):
    tmp = "%s.tmp.%d" % (path, os.getpid())
    with open(tmp, "wb") as fh:
        fh.write(data)
        fh.flush()
        os.fsync(fh.fileno())
    os.rename(tmp, path)


def _read_json(path):
    try:
        with open(path) as fh:
            return json.load(fh)
    except (OSError, ValueError):
        return None


def _loadavg():
    try:
        with open("/proc/loadavg") as fh:
            parts = fh.read().split()
        return [float(parts[0]), float(parts[1]), float(parts[2])]
    except (OSError, ValueError, IndexError):
        return None


def _meminfo():
    mem = {}
    try:
        with open("/proc/meminfo") as fh:
            for line in fh:
                if ":" not in line:
                    continue
                key, val = line.split(":", 1)
                val = val.strip().split()
                if val and val[0].isdigit():
                    mem[key.strip()] = int(val[0])
    except OSError:
        return None
    if not mem:
        return None
    return {"total_kb": mem.get("MemTotal"),
            "available_kb": mem.get("MemAvailable")}


def _workspace_free():
    try:
        st = os.statvfs(WORKSPACE)
        return st.f_bavail * st.f_frsize
    except OSError:
        return None


def _execd_task_count():
    data = _read_json(EXECD_TASKS)
    if not isinstance(data, dict):
        return None
    tasks = data.get("tasks")
    if not isinstance(tasks, list):
        return None
    running = [t for t in tasks
               if isinstance(t, dict) and t.get("state") == "running"]
    return {"total": len(tasks), "running": len(running)}


def _sync_summary():
    data = _read_json(SYNC_STATUS)
    if not isinstance(data, dict):
        return None
    return {"last": data.get("last"), "at": data.get("at"),
            "head": data.get("head"),
            "error": data.get("error")}


def build_heartbeat(pid, started_at, seq):
    """组装一次心跳 dict（纯函数，便于单测与 --once）。"""
    now = time.time()
    return {"service": "chv-daemon", "pid": pid,
            "started_at": started_at, "seq": seq,
            "at": now, "uptime_s": round(now - started_at, 1),
            "loadavg": _loadavg(), "mem": _meminfo(),
            "workspace_free_bytes": _workspace_free(),
            "execd_tasks": _execd_task_count(),
            "sync": _sync_summary()}


def _ask_execd_status(timeout=2.0):
    """问 execd 要任务状态；联系不上返回 (None, reason)。"""
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        sock.settimeout(timeout)
        sock.connect(EXECD_SOCK)
        sock.sendall(json.dumps({"op": "status"}).encode() + b"\n")
        buf = b""
        while len(buf) < MAX_REQUEST_BYTES:
            chunk = sock.recv(4096)
            if not chunk:
                break
            buf += chunk
            if buf.endswith(b"\n"):
                break
        data = json.loads(buf.decode())
        if not data.get("ok"):
            return None, "execd-status-not-ok"
        return data, None
    except (OSError, ValueError) as exc:
        return None, "execd-unreachable: %s" % exc
    finally:
        sock.close()


def prepare_snapshot():
    """快照配合：execd 有运行中任务或联系不上即拒绝。"""
    data, reason = _ask_execd_status()
    if data is None:
        return {"ok": False, "reason": reason}
    tasks = data.get("tasks") or []
    running = [t for t in tasks
               if isinstance(t, dict) and t.get("state") == "running"]
    if running:
        return {"ok": False, "reason": "task-running",
                "running": [t.get("id") for t in running]}
    marker = {"at": time.time(), "pid": os.getpid(),
              "note": "no running execd task"}
    _atomic_write(SNAPSHOT_READY_PATH,
                  (json.dumps(marker) + "\n").encode())
    return {"ok": True, "marker": SNAPSHOT_READY_PATH}


def handle_request(req):
    """socket 请求分发（纯函数，便于单测）。"""
    if not isinstance(req, dict):
        return {"ok": False, "error": "request must be an object"}
    op = req.get("op")
    if op == "ping":
        return {"ok": True, "service": "chv-daemon",
                "pid": os.getpid()}
    if op == "status":
        hb = _read_json(HEARTBEAT_PATH)
        return {"ok": True, "heartbeat": hb}
    if op == "prepare-snapshot":
        try:
            return prepare_snapshot()
        except OSError as exc:
            return {"ok": False, "reason": "marker-write-failed: %s"
                    % exc}
    return {"ok": False, "error": "unknown op: %s" % (op,)}


class Daemon:
    def __init__(self, interval=HEARTBEAT_INTERVAL):
        self.interval = interval
        self.pid = os.getpid()
        self.started_at = time.time()
        self.seq = 0
        old = _read_json(HEARTBEAT_PATH)
        if isinstance(old, dict) and isinstance(old.get("seq"), int):
            # 续号：宿主据 seq 连续判断心跳丢包还是 daemon 重启。
            self.seq = old["seq"]
        self._stop = threading.Event()
        self._server = None

    def write_heartbeat(self):
        self.seq += 1
        hb = build_heartbeat(self.pid, self.started_at, self.seq)
        _atomic_write(HEARTBEAT_PATH,
                      (json.dumps(hb) + "\n").encode())
        return hb

    def _serve(self):
        if os.path.exists(DAEMON_SOCK):
            os.remove(DAEMON_SOCK)
        srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        srv.bind(DAEMON_SOCK)
        os.chmod(DAEMON_SOCK, 0o600)
        srv.listen(8)
        srv.settimeout(0.5)
        self._server = srv
        while not self._stop.is_set():
            try:
                conn, _ = srv.accept()
            except socket.timeout:
                continue
            except OSError:
                break
            with conn:
                try:
                    conn.settimeout(5.0)
                    buf = b""
                    while len(buf) < MAX_REQUEST_BYTES:
                        chunk = conn.recv(4096)
                        if not chunk:
                            break
                        buf += chunk
                        if buf.endswith(b"\n"):
                            break
                    req = json.loads(buf.decode())
                    resp = handle_request(req)
                except (OSError, ValueError) as exc:
                    resp = {"ok": False, "error": str(exc)}
                try:
                    conn.sendall((json.dumps(resp) + "\n").encode())
                except OSError:
                    pass
        try:
            srv.close()
        except OSError:
            pass
        try:
            os.remove(DAEMON_SOCK)
        except OSError:
            pass

    def run(self):
        thread = threading.Thread(target=self._serve, daemon=True)
        thread.start()
        try:
            while not self._stop.is_set():
                self.write_heartbeat()
                self._stop.wait(self.interval)
        finally:
            self._stop.set()
            thread.join(timeout=5)

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


def main(argv=None):
    argv = list(sys.argv[1:] if argv is None else argv)
    interval = HEARTBEAT_INTERVAL
    once = False
    i = 0
    while i < len(argv):
        if argv[i] == "--interval" and i + 1 < len(argv):
            interval = float(argv[i + 1])
            i += 2
        elif argv[i] == "--once":
            once = True
            i += 1
        else:
            print("用法：chv-daemon [--interval 秒] [--once]",
                  file=sys.stderr)
            return 2
    daemon = Daemon(interval=interval)
    if once:
        daemon.write_heartbeat()
        print(json.dumps({"ok": True,
                          "heartbeat": HEARTBEAT_PATH,
                          "seq": daemon.seq}))
        return 0
    print("chv-daemon pid=%d heartbeat=%s" % (daemon.pid,
                                               HEARTBEAT_PATH),
          flush=True)

    def _on_signal(signum, _frame):
        daemon.stop()

    signal.signal(signal.SIGTERM, _on_signal)
    signal.signal(signal.SIGINT, _on_signal)
    daemon.run()
    return 0


if __name__ == "__main__":
    sys.exit(main())