#!/usr/bin/env python3
"""把一份固件推给板子（C2 OTA，firmware/main/ota.h）。

    python tools/ota_push.py b81f3fc9c8bc                 # 推 build/phonerail.bin，只下载不重启
    python tools/ota_push.py b81f3fc9c8bc --reboot        # 校验通过后立刻重启（⚠️ 五路全断一次）
    python tools/ota_push.py b81f3fc9c8bc --bin x.bin --no-watch
    python tools/ota_push.py b81f3fc9c8bc --status        # 只看板子 status 里的 ota 一块

做三件事：
  1. 算 sha256，把镜像传到服务器 /root/phonerail/fw/<sha256 前 16 位>.bin
     （nginx 7765 只放行这种文件名，rack/deploy/nginx-phonerail-fw.conf）；
  2. 在服务器上用后端账号经 127.0.0.1:1884 往 phonerail/<id>/cmd 发一条 `ota`；
  3. 盯着 rsp / tele / state 看进度，直到就绪 / 失败 / 新版本上线（--reboot 时）。

口令只在服务器上读（/root/phonerail/mqtt.env），不经过本机。
seq 用随机大数：板子的判重窗口是全局的，不能和 rack 后端的 seq 撞。
"""
import argparse
import hashlib
import json
import os
import random
import shlex
import subprocess
import sys

SSH_HOST = os.environ.get("PHONERAIL_SSH", "10.5.5.13")
FW_DIR = "/root/phonerail/fw"
PUBLIC_BASE = os.environ.get("PHONERAIL_FW_BASE", "http://home.example:7765/fw")
URL_MAX = 64  # ota.h OTA_URL_MAX
SLOT = 0x300000  # ota.h OTA_SLOT_BYTES

# 在服务器上跑：订阅 rsp/tele/state，发一条命令，逐行打印收到的消息（topic 前缀），超时退出。
REMOTE = r'''
set -e
. /root/phonerail/mqtt.env
T=phonerail/{id}
if [ "{send}" = 1 ]; then
  timeout {secs} mosquitto_sub -h 127.0.0.1 -p 1884 -u "$BACKEND_USER" -P "$BACKEND_PASS" -v \
      -t "$T/rsp" -t "$T/tele" -t "$T/state" &
  SUB=$!
  sleep 1
  mosquitto_pub -h 127.0.0.1 -p 1884 -u "$BACKEND_USER" -P "$BACKEND_PASS" -q 1 -t "$T/cmd" -m {msg}
  wait $SUB || true
else
  timeout {secs} mosquitto_sub -h 127.0.0.1 -p 1884 -u "$BACKEND_USER" -P "$BACKEND_PASS" -v \
      -t "$T/tele" -C 1 || true
fi
'''


def ssh(cmd: str, **kw):
    return subprocess.run(["ssh", SSH_HOST, cmd], **kw)


def watch(board: str, msg: str, secs: int, seq: int, reboot: bool, old_fw: str):
    """发命令并盯进度。返回退出码。"""
    script = REMOTE.format(id=board, send=1, secs=secs, msg=shlex.quote(msg))
    p = subprocess.Popen(["ssh", SSH_HOST, script], stdout=subprocess.PIPE, text=True,
                         encoding="utf-8", errors="replace")
    last = None
    rc = 1
    last_up = None
    rebooted = False
    acked = False     # 收到本条命令的回应了
    saw_dl = False    # 看到过「在下载」—— 之前那些 idle + err 是上一次留下的，不算这一次
    try:
        for line in p.stdout:
            topic, _, body = line.strip().partition(" ")
            try:
                j = json.loads(body)
            except ValueError:
                continue
            if topic.endswith("/rsp") and j.get("seq") == seq:
                print("回应：", body)
                if not j.get("ok"):
                    return 2
                acked = True
                if j.get("act") == "ready" and not reboot:
                    return 0
            elif topic.endswith("/tele"):
                up = j.get("up", 0)
                if last_up is not None and up < last_up:
                    rebooted = True
                    print("板子重启了（up %d → %d），fw %s" % (last_up, up, j.get("fw")))
                last_up = up
                o = j.get("ota") or {}
                cur = (o.get("st"), o.get("got"), o.get("err"))
                if cur != last:
                    last = cur
                    print("ota：", json.dumps(o, ensure_ascii=False), " fw", j.get("fw"))
                saw_dl = saw_dl or (acked and o.get("st") == "dl")
                if not acked:
                    continue
                if saw_dl and o.get("err") and o.get("st") == "idle":
                    print("失败：", o.get("err"))
                    return 3
                if saw_dl and o.get("st") == "ready" and not reboot:
                    print("已就绪：下次复位生效（要立刻换：再跑一次加 --reboot，不会重下）")
                    return 0
                if reboot and rebooted and o.get("pend") is False:
                    if o.get("bad") and j.get("fw") == old_fw:
                        print("⚠️ 回滚了：%s 被判无效，现在跑的是 %s" % (o.get("bad"), j.get("fw")))
                        return 4
                    print("固件 %s 已上线并自证健康" % j.get("fw"))
                    return 0
            elif topic.endswith("/state"):
                print("state：", body)
        print("超时（%d s）" % secs)
        return rc
    finally:
        p.kill()


def status(board: str) -> dict:
    out = ssh(REMOTE.format(id=board, send=0, secs=10, msg="''"), capture_output=True, text=True,
              encoding="utf-8").stdout
    _, _, body = out.strip().partition(" ")
    return json.loads(body) if body else {}


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("board", help="板子 ID（STA MAC 12 位小写十六进制）")
    ap.add_argument("--bin", default=os.path.join(os.path.dirname(__file__), "..", "build", "phonerail.bin"))
    ap.add_argument("--reboot", action="store_true", help="校验通过后立刻重启（五路全断一次）")
    ap.add_argument("--no-watch", action="store_true")
    ap.add_argument("--status", action="store_true")
    ap.add_argument("--timeout", type=int, default=300)
    a = ap.parse_args()

    st = status(a.board)
    if not st:
        print("10 s 内没收到这块板的遥测：板子不在线？")
        return 1
    print("板子当前 fw %s，ota %s" % (st.get("fw"), json.dumps(st.get("ota"), ensure_ascii=False)))
    if a.status:
        return 0

    data = open(a.bin, "rb").read()
    if not 0 < len(data) <= SLOT:
        print("镜像大小 %d 不在 (0, %d]" % (len(data), SLOT))
        return 1
    sha = hashlib.sha256(data).hexdigest()
    name = sha[:16] + ".bin"
    url = "%s/%s" % (PUBLIC_BASE, name)
    if len(url) > URL_MAX:
        print("URL 超过 %d 字节：%s" % (URL_MAX, url))
        return 1

    print("镜像 %s  %d 字节  sha256 %s" % (a.bin, len(data), sha))
    up = ssh("cat > %s/.%s.tmp && mv %s/.%s.tmp %s/%s && chmod 644 %s/%s && sha256sum %s/%s"
             % ((FW_DIR, name) * 5), input=data, capture_output=True)
    if up.returncode != 0 or sha not in up.stdout.decode():
        print("上传失败：", up.stderr.decode(errors="replace"))
        return 1
    print("已上传 →", url)

    seq = random.randint(1 << 30, (1 << 32) - 1)
    cmd = {"seq": seq, "cmd": "ota", "url": url, "sha256": sha, "size": len(data), "reboot": a.reboot}
    msg = json.dumps(cmd, separators=(",", ":"))
    if a.no_watch:
        ssh(REMOTE.format(id=a.board, send=1, secs=5, msg=shlex.quote(msg)))
        return 0
    return watch(a.board, msg, a.timeout, seq, a.reboot, st.get("fw"))


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