#!/usr/bin/env python3
"""客内 vsock notify 脚本（P1，feat-chv-muse-p1）。

启动后连宿主 CID 2 的 1024 端口发 ready，之后每 30 秒发一次
heartbeat。消息是 JSON 行：

    {"type": "ready"|"heartbeat", "vm": "<vm 名>", "ts": <unix 秒>}

宿主在 <vsock socket>_1024 这个 Unix socket 上监听接收
（cloud/chv/vsock.py 的 wait_ready/listen_notify）。

失败一律静默退出（exit 0，不打日志、不留 traceback）：vsock
只是就绪辅助信号，SSH 仍是唯一命令通道与就绪判定——本脚本
绝不能因为连不上、发失败而影响 sshd 或任务。调试时设
VSOCK_NOTIFY_DEBUG=1 可在 stderr 看到异常。

================ P2 镜像集成接入说明（本轨只写脚本，不改构建） ================
本文件不由 chv-guest/build.sh 安装，镜像集成是 P2 的范围。
P2 接入时做这三件：

1. 把本文件拷进镜像，例如 /app/vsock_notify.py（须 python3，
   客机镜像已有 Python；若最小镜像无 python3，P2 须先补）；
2. 加 systemd 单元，例如 /etc/systemd/system/vsock-notify.service：

       [Unit]
       Description=vsock notify to host (ready/heartbeat)
       After=multi-user.target
       [Service]
       Type=simple
       ExecStart=/usr/bin/python3 /app/vsock_notify.py %H
       Restart=no
       [Install]
       WantedBy=multi-user.target

   （%H 是客机 hostname；也可直接写 VM 名，或由 cloud-init
   写进环境变量 CHV_VM_NAME。Restart=no：本脚本失败即静默
   退出，不由 systemd 反复拉起打扰宿主监听。）
3. 前置条件（缺一则本脚本静默退出，任务仍走 SSH，不报错）：
   - 客机内核须含 CONFIG_VIRTIO_VSOCKETS；
   - 起 VM 时 CHV_VSOCK=1（默认开），VmConfig 带 vsock 设备
     （cloud/chv/executor.py 的 _balloon_vsock_payload）；
   - 宿主侧须在 VM boot 前后监听 <socket>_1024（否则客机
     首次 connect 失败，本脚本在 30 秒内重试，仍失败即退出）。

端口 1024 与宿主侧 cloud/chv/vsock.py 的 VSOCK_NOTIFY_PORT
须一致；改一处必须改另一处。
===============================================================================
"""

import json
import os
import socket
import sys
import time

# 宿主 CID（vsock 众所周知值，CHV 官方 docs/vsock.md）。
HOST_CID = 2
# notify 端口（与宿主侧 VSOCK_NOTIFY_PORT 一致）。
NOTIFY_PORT = 1024
# 首次 connect 重试窗口（秒）：宿主监听可能比客机启动晚到。
CONNECT_RETRY_SECONDS = 30
# 心跳间隔（秒）。
HEARTBEAT_INTERVAL = 30


def _vm_name(argv) -> str:
    """VM 名：argv[1] > CHV_VM_NAME > hostname。"""
    if len(argv) > 1 and argv[1]:
        return argv[1]
    env_name = os.environ.get("CHV_VM_NAME") or os.environ.get(
        "VSOCK_VM_NAME")
    if env_name:
        return env_name
    try:
        return socket.gethostname()
    except OSError:
        return "unknown"


def _send(sock, message_type: str, vm: str) -> None:
    line = json.dumps({"type": message_type, "vm": vm,
                       "ts": time.time()},
                      separators=(",", ":"))
    sock.sendall((line + "\n").encode("utf-8"))


def _connect(deadline: float):
    """在 deadline 前反复连 (HOST_CID, NOTIFY_PORT)，成功返回
    socket，超时返回 None。"""
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            return None
        sock = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM)
        try:
            sock.settimeout(min(remaining, 5))
            sock.connect((HOST_CID, NOTIFY_PORT))
            sock.settimeout(None)
            return sock
        except OSError as exc:
            try:
                sock.close()
            except OSError:
                pass
            if os.environ.get("VSOCK_NOTIFY_DEBUG") == "1":
                print("vsock_notify connect 失败：%r" % (exc,),
                      file=sys.stderr, flush=True)
            time.sleep(1)


def main(argv) -> int:
    vm = _vm_name(argv)
    deadline = time.monotonic() + CONNECT_RETRY_SECONDS
    sock = _connect(deadline)
    if sock is None:
        return 0
    try:
        _send(sock, "ready", vm)
        while True:
            time.sleep(HEARTBEAT_INTERVAL)
            _send(sock, "heartbeat", vm)
    except (OSError, ValueError):
        # 任何发送失败都静默退出，不影响 sshd/任务。
        return 0
    except KeyboardInterrupt:
        return 0
    finally:
        try:
            sock.close()
        except OSError:
            pass
    return 0


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