#!/usr/bin/env python3
"""PhoneRail USB 命令链路的 PC 端客户端（A13，需求 §10.8 / firmware/main/ulink.h）。

    python prlink.py COM5 status
    python prlink.py COM5 '{"cmd":"key","rail":3,"keys":"pwr","ms":100}'
    python prlink.py COM5 status --logs          # 同时把板子日志打到 stderr
    python prlink.py --selftest                  # 只跑本地的 CRC / 解帧自检，不碰串口

线上格式（与 ulink.h 一致，改一边必须改另一边）：
    发  @C <json>*XXXX\\n      XXXX = CRC-16/CCITT-FALSE(json)，4 位大写十六进制
    收  @R <json>*XXXX\\n      CRC 必带；对不上的整行丢弃
    其余任何一行都是日志。

两条上位机侧的规矩（cmd.h）：
  · seq 由这里生成，**重传用同一个 seq**：超时后原样再发，板子判重兜住，
    不会把「按 10 秒电源键」执行两遍；想「再试一次」就是一条新命令、新 seq。
  · 打开串口时 **DTR / RTS 都保持低**。USB-Serial-JTAG 会把特定的 DTR/RTS 组合
    当成复位 / 进下载模式的信号（esptool 就是这么做的），默认打开方式可能把板子复位。
"""
import argparse
import json
import sys
import time

CMD_PREFIX = "@C "
RSP_PREFIX = "@R "


def crc16(data: bytes) -> int:
    crc = 0xFFFF
    for b in data:
        crc ^= b << 8
        for _ in range(8):
            crc = ((crc << 1) ^ 0x1021) if (crc & 0x8000) else (crc << 1)
            crc &= 0xFFFF
    return crc


def frame(obj: dict) -> bytes:
    js = json.dumps(obj, separators=(",", ":"), ensure_ascii=True)
    return f"{CMD_PREFIX}{js}*{crc16(js.encode()):04X}\n".encode()


def deframe(line: str):
    """一行输出 → 回应 dict；日志 / CRC 错 / 不带 CRC 都返回 None（同 ulink_deframe）。"""
    line = line.rstrip("\r")
    if not line.startswith(RSP_PREFIX) or len(line) < len(RSP_PREFIX) + 5:
        return None
    body = line[len(RSP_PREFIX):]
    js, star, tail = body[:-5], body[-5], body[-4:]
    if star != "*":
        return None
    try:
        want = int(tail, 16)
    except ValueError:
        return None
    if crc16(js.encode()) != want:
        return None
    try:
        return json.loads(js)
    except ValueError:
        return None


def new_seq() -> int:
    """跨进程也不重复的 seq：0.1 s 分辨率的时间戳，落在 1~2^32−1。
    板子的判重窗口按「值相等」判（cmd.h），不要求单调，只要求不撞。"""
    return int(time.time() * 10) % 0xFFFFFFFF + 1


class Link:
    def __init__(self, port: str, show_logs: bool = False):
        import serial  # 只有真要连串口时才需要 pyserial

        self.ser = serial.Serial()
        self.ser.port = port
        self.ser.baudrate = 115200  # USB-Serial-JTAG 不看波特率，填什么都行
        self.ser.timeout = 0.05
        self.ser.dtr = False  # 见文件头：先置低再 open
        self.ser.rts = False
        self.ser.open()
        self.show_logs = show_logs
        self.buf = b""

    def close(self):
        self.ser.close()

    def _lines(self, deadline: float):
        while time.monotonic() < deadline:
            # 有多少读多少；什么都没有时读 1 字节、最多等 timeout。
            # 写成 read(512) 的话，凑不满 512 字节每次都要白等满 timeout。
            chunk = self.ser.read(self.ser.in_waiting or 1)
            if chunk:
                self.buf += chunk
            while b"\n" in self.buf:
                raw, self.buf = self.buf.split(b"\n", 1)
                yield raw.decode("utf-8", errors="replace")

    def call(self, obj: dict, timeout: float = 1.0, retries: int = 3) -> dict:
        obj = dict(obj)
        if "seq" not in obj:
            obj["seq"] = new_seq()
        data = frame(obj)
        for _ in range(retries):
            self.ser.write(data)
            for line in self._lines(time.monotonic() + timeout):
                rsp = deframe(line)
                if rsp is None:
                    if self.show_logs and line.strip():
                        print(line, file=sys.stderr)
                    continue
                if rsp.get("seq") == obj["seq"]:
                    return rsp
                # 别的 seq 的回应（上一条命令迟到的那一份）：丢掉接着等
        raise TimeoutError(f"seq {obj['seq']} 等了 {retries} 次都没有回应")


def selftest() -> int:
    assert crc16(b"123456789") == 0x29B1
    js = '{"seq":7,"ok":true}'
    good = f"@R {js}*{crc16(js.encode()):04X}"
    assert deframe(good) == {"seq": 7, "ok": True}
    assert deframe(good + "\r") == {"seq": 7, "ok": True}
    assert deframe(good[:-1] + ("0" if good[-1] != "0" else "1")) is None  # CRC 错
    assert deframe(f"@R {js}") is None  # 没 CRC
    assert deframe(f" {good}") is None  # 不在行首（日志里垫过空格的伪帧）
    assert deframe("I (1) x: hello") is None
    f = frame({"cmd": "status", "seq": 1}).decode()
    assert f.startswith('@C {"cmd":"status","seq":1}*') and f.endswith("\n")
    print("prlink 自检通过")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("port", nargs="?")
    ap.add_argument("command", nargs="?", help='命令名（如 status）或一整段 JSON')
    ap.add_argument("--logs", action="store_true", help="把板子日志打到 stderr")
    ap.add_argument("--timeout", type=float, default=1.0)
    ap.add_argument("--selftest", action="store_true")
    a = ap.parse_args()
    if a.selftest:
        return selftest()
    if not a.port or not a.command:
        ap.error("需要 port 与 command")
    obj = json.loads(a.command) if a.command.lstrip().startswith("{") else {"cmd": a.command}
    link = Link(a.port, show_logs=a.logs)
    try:
        rsp = link.call(obj, timeout=a.timeout)
    finally:
        link.close()
    print(json.dumps(rsp, ensure_ascii=False))
    return 0 if rsp.get("ok") else 1


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