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

职责（只做这三件）：
  1. shell 命令：op exec（同步跑完返回）/ spawn（脱离式后台跑，
     立即返回 id）。一律 list argv、绝不用 shell=True；工作区
     内命令以 chv 身份跑（runuser -u chv），codex 除外。
  2. 文件读写：op read/write，路径只允许 /home/chv/workspace
     之内（realpath + commonpath 校验），越界即拒绝。
  3. codex 进程托管：op codex——经 chv-cell run 起
     `python3 /app/host.py <args>`（单元隔离口径不变，不绕
     cell 直跑 host.py）；子进程 setsid 脱离，execd 崩溃不
     带走它；任务登记落盘 /run/chv/execd-tasks.json，execd
     重启后按 pid 重新认领（reattach），任务状态不丢。

客户端鉴权：SO_PEERCRED 的 uid 须是 0 或 1000（root / chv），
其余一律拒绝。socket /run/chv/execd.sock 0600 root:root。
调用方是宿主执行器经 SSH（root）或客机内 root 服务；单元内
codex 不直接调 execd（走 chv-broker 白名单）。

任务状态机：running → completed（rc 已收）/ failed（spawn
失败）/ killed / unknown（execd 重启后 pid 已死且退出码
不可考——init 已收走，如实记 unknown，不编 rc）。
spawn 输出进 /run/chv/tasks/<id>.log（runuser 以 chv 写，
codex 以 root 写，均 0644）。

只用标准库。以 root 运行（要起 chv-cell + runuser）；单租户
灰度口径——与今天“执行器经 SSH 以 root 跑 host.py”同权，
多租户收紧待安全再审翻绿后另议。
"""
import json
import os
import signal
import socket
import struct
import subprocess
import sys
import threading
import time
import uuid

EXECD_SOCK = "/run/chv/execd.sock"
TASKS_PATH = "/run/chv/execd-tasks.json"
TASKS_LOG_DIR = "/run/chv/tasks"
WORKSPACE = "/home/chv/workspace"
HOST_PY = "/app/host.py"
CELL_BIN = "/usr/local/bin/chv-cell"
ALLOWED_UIDS = (0, 1000)
MAX_REQUEST_BYTES = 65536
MAX_OUTPUT_BYTES = 65536
MAX_FILE_BYTES = 1 << 20
DEFAULT_TIMEOUT = 120.0

# codex 子进程可透传的环境变量白名单（前缀匹配；调用方
# 传什么别的都丢弃。模型网关 key 走这个口径，不走 argv，
# 绝不进任务登记文件与日志）。
CODEX_ENV_PREFIXES = ("OPENAI_", "CODEX_", "CHV_MODEL_",
                      "PLAYWRIGHT_")


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 _peer_uid(conn):
    """SO_PEERCRED 取对端 uid；取不到即拒绝（不降级）。"""
    try:
        creds = conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED,
                                struct.calcsize("3i"))
        _pid, uid, _gid = struct.unpack("3i", creds)
        return uid
    except (OSError, struct.error):
        return None


def _within_workspace(path):
    """路径是否在 WORKSPACE 内（realpath 后 commonpath 校验）。"""
    try:
        real = os.path.realpath(path)
    except OSError:
        return False
    return os.path.commonpath([real, WORKSPACE]) == WORKSPACE


def _truncate(text, limit=MAX_OUTPUT_BYTES):
    if len(text) > limit:
        return text[:limit] + "\n...<truncated %d bytes>" % (
            len(text) - limit)
    return text


def _codex_env(extra):
    """调用方 env 过白名单；只留前缀命中的项。"""
    env = {"PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin"),
           "HOME": "/home/chv",
           "PLAYWRIGHT_BROWSERS_PATH":
               "/usr/local/share/ms-playwright"}
    if isinstance(extra, dict):
        for key, val in extra.items():
            if isinstance(key, str) and isinstance(val, str) \
                    and key.startswith(CODEX_ENV_PREFIXES):
                env[key] = val
    return env


class TaskStore:
    """任务登记（落盘 JSON，execd 重启后重读认领）。"""

    def __init__(self, path=TASKS_PATH):
        self.path = path
        self._lock = threading.Lock()
        self.tasks = {}
        data = _read_json(path)
        if isinstance(data, dict) and isinstance(
                data.get("tasks"), list):
            for item in data["tasks"]:
                if isinstance(item, dict) and item.get("id"):
                    self.tasks[item["id"]] = item
        self._reattach()

    def _reattach(self):
        """重启后认领：running 任务按 pid 查存活；已死记
        unknown（退出码不可考，不编），活着的保持 running。"""
        for task in self.tasks.values():
            if task.get("state") != "running":
                continue
            pid = task.get("pid")
            alive = False
            if isinstance(pid, int):
                try:
                    os.kill(pid, 0)
                    alive = True
                except (OSError, ProcessLookupError):
                    alive = False
                except PermissionError:
                    alive = True
            if alive:
                task["note"] = "reattached after execd restart"
            else:
                task["state"] = "unknown"
                task["note"] = ("execd restarted, exit code "
                                "unrecoverable (reaped by init)")
        self._save_locked()

    def _save_locked(self):
        data = {"tasks": [self.tasks[k] for k in sorted(self.tasks)]}
        _atomic_write(self.path,
                      (json.dumps(data) + "\n").encode())

    def add(self, task):
        with self._lock:
            self.tasks[task["id"]] = task
            self._save_locked()
        return task

    def update(self, task_id, **fields):
        with self._lock:
            task = self.tasks.get(task_id)
            if task is None:
                return None
            task.update(fields)
            self._save_locked()
            return task

    def get(self, task_id):
        with self._lock:
            task = self.tasks.get(task_id)
            return dict(task) if task else None

    def list(self):
        with self._lock:
            return [dict(self.tasks[k]) for k in sorted(self.tasks)]


class Execd:
    def __init__(self, tasks_path=TASKS_PATH):
        self.store = TaskStore(tasks_path)
        self._stop = threading.Event()
        self._procs = {}
        self._procs_lock = threading.Lock()
        self._reaper = threading.Thread(target=self._reap_loop,
                                        daemon=True)

    def _register_proc(self, task_id, proc):
        with self._procs_lock:
            self._procs[task_id] = proc

    def _reap_loop(self):
        """收割已退出的子进程，登记 rc；未知 pid 的直接丢弃。
        重启后认领的 running 任务不在 _procs 里：轮询 pid
        存活，死了记 unknown（退出码不可考，不编）。"""
        while not self._stop.is_set():
            with self._procs_lock:
                items = list(self._procs.items())
            tracked = set()
            for task_id, proc in items:
                tracked.add(task_id)
                rc = proc.poll()
                if rc is None:
                    continue
                with self._procs_lock:
                    self._procs.pop(task_id, None)
                task = self.store.get(task_id)
                if task and task.get("state") == "running":
                    self.store.update(task_id, state="completed",
                                      rc=rc, ended_at=time.time())
            for task in self.store.list():
                if task.get("state") != "running" \
                        or task.get("id") in tracked:
                    continue
                pid = task.get("pid")
                alive = False
                if isinstance(pid, int):
                    try:
                        os.kill(pid, 0)
                        alive = True
                    except OSError:
                        alive = False
                if not alive:
                    self.store.update(
                        task["id"], state="unknown",
                        note=("execd restarted, exit code "
                              "unrecoverable (reaped by init)"),
                        ended_at=time.time())
            self._stop.wait(1.0)

    # -- op 实现 ---------------------------------------------------------
    def op_exec(self, req):
        argv = req.get("argv")
        if not isinstance(argv, list) or not argv \
                or not all(isinstance(a, str) for a in argv):
            return {"ok": False, "error": "argv must be a non-empty list of strings"}
        cwd = req.get("cwd") or WORKSPACE
        if not _within_workspace(cwd) and cwd != "/home/chv":
            return {"ok": False, "error": "cwd outside workspace: %s" % cwd}
        try:
            timeout = float(req.get("timeout") or DEFAULT_TIMEOUT)
        except (TypeError, ValueError):
            return {"ok": False, "error": "bad timeout"}
        timeout = min(max(timeout, 1.0), 600.0)
        cmd = ["runuser", "-u", "chv", "--"] + argv
        try:
            proc = subprocess.run(cmd, cwd=cwd, timeout=timeout,
                                  capture_output=True, text=True,
                                  errors="replace")
        except subprocess.TimeoutExpired as exc:
            return {"ok": False, "error": "timeout after %.0fs" % timeout,
                    "stdout": _truncate(exc.stdout or ""),
                    "stderr": _truncate(exc.stderr or "")}
        except OSError as exc:
            return {"ok": False, "error": "spawn failed: %s" % exc}
        return {"ok": True, "rc": proc.returncode,
                "stdout": _truncate(proc.stdout or ""),
                "stderr": _truncate(proc.stderr or "")}

    def op_spawn(self, req, kind="exec"):
        argv = req.get("argv")
        if not isinstance(argv, list) or not argv \
                or not all(isinstance(a, str) for a in argv):
            return {"ok": False, "error": "argv must be a non-empty list of strings"}
        cwd = req.get("cwd") or WORKSPACE
        if not _within_workspace(cwd) and cwd != "/home/chv":
            return {"ok": False, "error": "cwd outside workspace: %s" % cwd}
        if kind == "codex":
            if not os.path.exists(CELL_BIN):
                return {"ok": False, "error": "chv-cell not found"}
            full = [CELL_BIN, "run", "--", "python3", HOST_PY] + argv
            env = _codex_env(req.get("env"))
            as_user = None
        else:
            full = ["runuser", "-u", "chv", "--"] + argv
            env = None
            as_user = "chv"
        task_id = uuid.uuid4().hex[:12]
        os.makedirs(TASKS_LOG_DIR, exist_ok=True)
        log_path = os.path.join(TASKS_LOG_DIR, task_id + ".log")
        try:
            log_fh = open(log_path, "ab", buffering=0)
        except OSError as exc:
            return {"ok": False, "error": "log open failed: %s" % exc}
        try:
            proc = subprocess.Popen(
                full, cwd=cwd, stdout=log_fh, stderr=log_fh,
                stdin=subprocess.DEVNULL, start_new_session=True,
                env=env, close_fds=True)
        except OSError as exc:
            log_fh.close()
            return {"ok": False, "error": "spawn failed: %s" % exc}
        task = {"id": task_id, "kind": kind,
                "argv_preview": argv[:4], "pid": proc.pid,
                "state": "running", "started_at": time.time(),
                "log": log_path,
                "as_user": as_user or "root-via-cell"}
        self.store.add(task)
        self._register_proc(task_id, proc)
        return {"ok": True, "id": task_id, "pid": proc.pid,
                "log": log_path}

    def op_read(self, req):
        path = req.get("path")
        if not isinstance(path, str) or not path:
            return {"ok": False, "error": "path required"}
        if not _within_workspace(path):
            return {"ok": False, "error": "path outside workspace"}
        try:
            with open(path, "rb") as fh:
                data = fh.read(MAX_FILE_BYTES + 1)
        except OSError as exc:
            return {"ok": False, "error": "read failed: %s" % exc}
        truncated = len(data) > MAX_FILE_BYTES
        return {"ok": True, "path": path,
                "content": data[:MAX_FILE_BYTES].decode(
                    "utf-8", errors="replace"),
                "truncated": truncated}

    def op_write(self, req):
        path = req.get("path")
        content = req.get("content")
        if not isinstance(path, str) or not path:
            return {"ok": False, "error": "path required"}
        if not isinstance(content, str):
            return {"ok": False, "error": "content must be a string"}
        if not _within_workspace(path):
            return {"ok": False, "error": "path outside workspace"}
        try:
            parent = os.path.dirname(path)
            if parent:
                os.makedirs(parent, exist_ok=True)
            with open(path, "w") as fh:
                fh.write(content)
        except OSError as exc:
            return {"ok": False, "error": "write failed: %s" % exc}
        return {"ok": True, "path": path, "bytes": len(content)}

    def op_task_status(self, req):
        task_id = req.get("id")
        if task_id:
            task = self.store.get(task_id)
            if task is None:
                return {"ok": False, "error": "no such task: %s" % task_id}
            return {"ok": True, "task": task}
        return {"ok": True, "tasks": self.store.list()}

    def op_task_kill(self, req):
        task_id = req.get("id")
        if not task_id:
            return {"ok": False, "error": "id required"}
        task = self.store.get(task_id)
        if task is None:
            return {"ok": False, "error": "no such task: %s" % task_id}
        if task.get("state") != "running":
            return {"ok": False,
                    "error": "task not running: %s" % task.get("state")}
        pid = task.get("pid")
        sig = signal.SIGTERM
        if req.get("signal") == "KILL":
            sig = signal.SIGKILL
        try:
            os.kill(pid, sig)
        except OSError as exc:
            return {"ok": False, "error": "kill failed: %s" % exc}
        self.store.update(task_id, state="killed",
                          ended_at=time.time())
        return {"ok": True, "id": task_id, "pid": pid}

    def handle(self, req):
        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-execd",
                    "pid": os.getpid()}
        if op == "status":
            return {"ok": True, "tasks": self.store.list()}
        if op == "exec":
            return self.op_exec(req)
        if op == "spawn":
            return self.op_spawn(req, kind="exec")
        if op == "codex":
            return self.op_spawn(req, kind="codex")
        if op == "read":
            return self.op_read(req)
        if op == "write":
            return self.op_write(req)
        if op == "task-status":
            return self.op_task_status(req)
        if op == "task-kill":
            return self.op_task_kill(req)
        return {"ok": False, "error": "unknown op: %s" % (op,)}

    # -- socket 服务 ------------------------------------------------------
    def serve(self):
        if os.path.exists(EXECD_SOCK):
            os.remove(EXECD_SOCK)
        srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        srv.bind(EXECD_SOCK)
        os.chmod(EXECD_SOCK, 0o600)
        srv.listen(8)
        srv.settimeout(0.5)
        self._reaper.start()
        print("chv-execd pid=%d sock=%s" % (os.getpid(), EXECD_SOCK),
              flush=True)
        try:
            while not self._stop.is_set():
                try:
                    conn, _ = srv.accept()
                except socket.timeout:
                    continue
                except OSError:
                    break
                with conn:
                    uid = _peer_uid(conn)
                    if uid not in ALLOWED_UIDS:
                        resp = {"ok": False,
                                "error": "uid %s not allowed" % (uid,)}
                        try:
                            conn.sendall(
                                (json.dumps(resp) + "\n").encode())
                        except OSError:
                            pass
                        continue
                    try:
                        conn.settimeout(10.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 = self.handle(req)
                    except (OSError, ValueError) as exc:
                        resp = {"ok": False, "error": str(exc)}
                    try:
                        conn.sendall((json.dumps(resp) + "\n").encode())
                    except OSError:
                        pass
        finally:
            self._stop.set()
            try:
                srv.close()
            except OSError:
                pass
            try:
                os.remove(EXECD_SOCK)
            except OSError:
                pass

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


def main(argv=None):
    argv = list(sys.argv[1:] if argv is None else argv)
    if argv:
        print("用法：chv-execd（无参数）", file=sys.stderr)
        return 2
    execd = Execd()

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

    signal.signal(signal.SIGTERM, _on_signal)
    signal.signal(signal.SIGINT, _on_signal)
    execd.serve()
    return 0


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