#!/usr/bin/env python3
"""通过电脑串口测试STM8 Mesh继电器的主程序OTA功能。

发送流程：
    0x08：发送固件长度和版本
    0x09：逐包发送固件数据，并等待设备返回下一期待偏移
    0x0A：发送完整固件CRC16
    0x0B：请求设备分片校验W25Q80BV中的固件
    0x0D：设备复位并升级后主动上报升级结果

当前协议约定：
    - JDY-28M发送外层：41 54 2B AA A1 | DST_H DST_L | A0 | BUSINESS | 0D 0A
    - 业务帧：CMD | LEN | DATA | CRC16_H CRC16_L
    - CRC算法：CRC-16/MODBUS
    - DATA中的多字节数值和CRC均为高字节在前
    - CMD=0x09每帧最多携带2字节固件数据

依赖：
    python3 -m pip install pyserial

示例：
    python3 tools/ota_serial_test.py \
        --port /dev/cu.usbserial-0001 \
        --firmware build/mesh_relay.bin \
        --source 0001 \
        --target 000C
"""

from __future__ import annotations

import argparse
import re
import struct
import sys
import time
from pathlib import Path

try:
    import serial
except ImportError as exc:
    raise SystemExit(
        "缺少pyserial，请先执行：python3 -m pip install pyserial"
    ) from exc


# JDY-28M Mesh发送外层固定字段。
MESH_TX_PREFIX = bytes.fromhex("41 54 2B AA A1")
MESH_DATA_FLAG = 0xA0
MESH_TAIL = bytes.fromhex("0A")

# JDY-28M接收数据时，串口侧看到的固定前缀。
MESH_RX_PREFIX = bytes.fromhex("F1 DD 09")

# OTA命令定义。
CMD_OTA_START = 0x08
CMD_OTA_DATA = 0x09
CMD_OTA_FINISH = 0x0A
CMD_OTA_VERIFY = 0x0B
CMD_OTA_RESULT_REPORT = 0x0D
CMD_OTA_RESULT_ACK = 0x8D
CMD_FIRMWARE_ANNOUNCE = 0x10
CMD_FIRMWARE_ANNOUNCE_RESPONSE = 0x90

# 协议和目标MCU的限制。
MAX_FIRMWARE_SIZE = 0x2000
OTA_CHUNK_SIZE = 2


def crc16_modbus(data: bytes) -> int:
    """计算CRC-16/MODBUS。

    CRC计算范围不包含Mesh外层帧头、目标地址、A0和帧尾，只包含业务帧：
    SRC、CMD、LEN和DATA。
    """

    crc = 0xFFFF
    for value in data:
        crc ^= value
        for _ in range(8):
            crc = (crc >> 1) ^ 0xA001 if crc & 1 else crc >> 1
    return crc & 0xFFFF


def parse_hex_address(value: str, name: str) -> int:
    """解析两字节十六进制Mesh短地址。"""

    try:
        address = int(value, 16)
    except ValueError as exc:
        raise argparse.ArgumentTypeError(f"{name}不是有效的十六进制地址：{value}") from exc
    if not 0 <= address <= 0xFFFF:
        raise argparse.ArgumentTypeError(f"{name}必须在0000～FFFF范围内")
    return address


def read_firmware_version(header_path: Path) -> int:
    """从主程序app_config.h读取FIRMWARE_VERSION。"""

    content = header_path.read_text(encoding="utf-8")
    match = re.search(
        r"^\s*#define\s+FIRMWARE_VERSION\s+((?:0[xX][0-9A-Fa-f]+)|(?:[0-9]+))",
        content,
        flags=re.MULTILINE,
    )
    if match is None:
        raise ValueError(f"在{header_path}中找不到FIRMWARE_VERSION")
    version = int(match.group(1), 0)
    if not 0 <= version <= 0xFFFF:
        raise ValueError(f"固件版本超出16位范围：{version}")
    return version


def make_business_frame(source: int, command: int, data: bytes) -> bytes:
    """生成业务帧，末尾CRC16按高字节在前发送。"""

    if len(data) > 4:
        raise ValueError("当前JDY-28M外层20字节限制下，业务DATA不能超过4字节")
    # JDY-28M外层已经携带发送方和接收方地址，业务帧不再重复嵌入source。
    body = struct.pack(">BB", command, len(data)) + data
    crc = crc16_modbus(body)
    return body + struct.pack(">H", crc)


def make_mesh_frame(target: int, business: bytes) -> bytes:
    """封装JDY-28M Mesh发送帧。"""

    return (
        MESH_TX_PREFIX
        + struct.pack(">H", target)
        + bytes([MESH_DATA_FLAG])
        + business
        + MESH_TAIL
    )


def send_command(ser: serial.Serial, source: int, target: int,
                 command: int, data: bytes) -> None:
    """通过串口发送一条完整Mesh业务命令。"""

    frame = make_mesh_frame(target, make_business_frame(source, command, data))
    ser.write(frame)
    ser.flush()


def parse_received_business(frame: bytes) -> tuple[int, int, bytes] | None:
    """从JDY-28M接收外层帧中提取业务命令并校验CRC。

    接收格式为：
        F1 DD 09 | SRC地址2字节 | DST地址2字节 | BUSINESS | 0A

    发送方和接收方地址由JDY-28M外层提供，业务帧不再重复携带地址。
    """

    if len(frame) < len(MESH_RX_PREFIX) + 4 + 4 + 1:
        return None
    if not frame.startswith(MESH_RX_PREFIX) or not frame.endswith(MESH_TAIL):
        return None

    source = struct.unpack(">H", frame[3:5])[0]
    business = frame[len(MESH_RX_PREFIX) + 4:-1]
    if len(business) < 4:
        return None
    command, length = struct.unpack(">BB", business[:2])
    if len(business) != 2 + length + 2 or length > 4:
        return None
    body = business[:2 + length]
    received_crc = struct.unpack(">H", business[2 + length:])[0]
    if crc16_modbus(body) != received_crc:
        return None
    return source, command, business[2:2 + length]


def receive_response(ser: serial.Serial, timeout: float) -> tuple[int, int, bytes]:
    """等待并解析一条合法的JDY-28M业务响应。"""

    deadline = time.monotonic() + timeout
    buffer = bytearray()
    while time.monotonic() < deadline:
        waiting = ser.in_waiting
        if waiting:
            buffer.extend(ser.read(waiting))
        else:
            time.sleep(0.005)

        while True:
            start = buffer.find(MESH_RX_PREFIX)
            if start < 0:
                # 保留可能是固定前缀开头的最后几个字节，避免跨次读取丢帧。
                keep = min(len(buffer), len(MESH_RX_PREFIX) - 1)
                if keep:
                    del buffer[:-keep]
                else:
                    buffer.clear()
                break
            if start:
                del buffer[:start]
            if len(buffer) < len(MESH_RX_PREFIX) + 4 + 4 + 1:
                break

            data_length = buffer[len(MESH_RX_PREFIX) + 4 + 1]
            total = len(MESH_RX_PREFIX) + 4 + 2 + data_length + 2 + 1
            if data_length > 4:
                del buffer[0]
                continue
            if len(buffer) < total:
                break

            candidate = bytes(buffer[:total])
            del buffer[:total]
            parsed = parse_received_business(candidate)
            if parsed is not None:
                return parsed

    raise TimeoutError(f"等待设备响应超时（{timeout:.1f}秒）")


def wait_command_response(ser: serial.Serial, expected_command: int,
                          timeout: float) -> bytes:
    """等待指定应答命令，忽略上线通知等无关业务帧。"""

    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        source, command, data = receive_response(
            ser, max(0.1, deadline - time.monotonic())
        )
        print(f"收到响应：SRC={source:04X} CMD={command:02X} DATA={data.hex(' ').upper()}")
        if command == expected_command:
            return data
    raise TimeoutError(f"未收到期望响应CMD={expected_command:02X}")


def wait_upgrade_result(ser: serial.Serial, timeout: float) -> tuple[int, int]:
    """等待设备复位后主动上报CMD=0x0D升级结果。

    0x8B只表示外部Flash固件校验成功。设备随后进入Bootloader搬运固件，
    重启主程序并在蓝牙初始化完成后广播0x0D。因此这里除了结果码，还要
    保留JDY-28M外层提供的SRC地址，以便后续发送0x8D确认。

    @return (设备短地址, 升级结果码)。
    @raises TimeoutError 在指定时间内未收到合法0x0D。
    """
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            source, command, data = receive_response(
                ser, max(0.1, deadline - time.monotonic())
            )
        except TimeoutError:
            break
        print(f"收到升级结果：SRC={source:04X} CMD={command:02X} "
              f"DATA={data.hex(' ').upper()}")
        if command == CMD_OTA_RESULT_REPORT and len(data) == 1:
            return source, data[0]
    raise TimeoutError(f"等待设备升级结果超时（{timeout:.1f}秒）")


def send_ota(args: argparse.Namespace) -> None:
    """执行完整OTA升级测试。"""

    firmware = Path(args.firmware).read_bytes()
    if not firmware:
        raise ValueError("固件文件为空")
    if len(firmware) > MAX_FIRMWARE_SIZE:
        raise ValueError(f"固件大小{len(firmware)}超过8KB限制")

    version_header = Path(args.version_header)
    version = args.version if args.version is not None else read_firmware_version(version_header)
    crc = crc16_modbus(firmware)
    print(f"固件：{args.firmware}")
    print(f"长度：{len(firmware)} 字节，版本：0x{version:04X}，CRC16：0x{crc:04X}")

    with serial.Serial(args.port, args.baudrate, timeout=0) as ser:
        # 进入OTA：长度和版本均为高字节在前。
        send_command(ser, args.source, args.target, CMD_OTA_START,
                     struct.pack(">HH", len(firmware), version))
        status = wait_command_response(ser, 0x88, args.timeout)
        if status != b"\x00":
            raise RuntimeError(f"设备拒绝进入OTA，状态={status.hex(' ').upper()}")

        offset = 0
        while offset < len(firmware):
            payload = firmware[offset:offset + OTA_CHUNK_SIZE]
            data = struct.pack(">H", offset) + payload
            send_command(ser, args.source, args.target, CMD_OTA_DATA, data)
            response = wait_command_response(ser, 0x89, args.timeout)
            if len(response) != 3 or response[0] != 0:
                raise RuntimeError(f"分包写入失败，响应={response.hex(' ').upper()}")
            expected_offset = struct.unpack(">H", response[1:3])[0]
            expected_local = offset + len(payload)
            if expected_offset != expected_local:
                raise RuntimeError(
                    f"设备期待偏移0x{expected_offset:04X}，本地下一偏移为0x{expected_local:04X}"
                )
            offset = expected_offset
            print(f"OTA分包进度：{offset}/{len(firmware)}", end="\r", flush=True)
        print()

        # 所有分包均已确认后，再发送完整固件CRC16。
        send_command(ser, args.source, args.target, CMD_OTA_FINISH,
                     struct.pack(">H", crc))
        status = wait_command_response(ser, 0x8A, args.timeout)
        if status != b"\x00":
            raise RuntimeError(f"设备拒绝固件CRC，状态={status.hex(' ').upper()}")

        # 请求设备后台分片读取W25Q80BV并计算完整CRC。
        send_command(ser, args.source, args.target, CMD_OTA_VERIFY, b"")
        status = wait_command_response(ser, 0x8B, args.verify_timeout)
        if status != b"\x00":
            raise RuntimeError(f"设备固件校验失败，状态={status.hex(' ').upper()}")

        # 0x8B只表示外部Flash中的固件校验成功。设备随后自动复位，
        # 由Bootloader搬运主程序；这里继续等待主程序启动后的0x0D上报。
        device_address, upgrade_result = wait_upgrade_result(
            ser, args.result_timeout
        )
        if upgrade_result != 0x01:
            raise RuntimeError(
                f"设备升级失败：设备地址={device_address:04X}，"
                f"结果码=0x{upgrade_result:02X}"
            )

        # 0x01表示Bootloader搬运并校验成功。回复0x8D后，设备清除待上报
        # 结果，避免后续每次启动都重复上报同一条升级结果。
        send_command(ser, args.source, device_address,
                     CMD_OTA_RESULT_ACK, bytes([upgrade_result]))

    print("OTA升级成功：设备已复位、Bootloader已搬运主程序，并完成结果确认。")


def discover_upgrade_devices(args: argparse.Namespace) -> None:
    """广播新固件信息并收集需要升级的设备列表。"""

    firmware = Path(args.firmware).read_bytes()
    if not firmware or len(firmware) > MAX_FIRMWARE_SIZE:
        raise ValueError("固件为空或超过8KB限制")
    version = args.version if args.version is not None else read_firmware_version(
        Path(args.version_header)
    )

    print(f"广播固件版本：0x{version:04X}，长度：{len(firmware)}字节")
    devices: dict[int, tuple[int, int]] = {}
    with serial.Serial(args.port, args.baudrate, timeout=0) as ser:
        # 0x10使用广播地址FFFF，仅通知版本和长度，不会直接启动OTA。
        announce = struct.pack(">HH", version, len(firmware))
        send_command(ser, args.source, 0xFFFF,
                     CMD_FIRMWARE_ANNOUNCE, announce)

        deadline = time.monotonic() + args.discovery_timeout
        while time.monotonic() < deadline:
            try:
                source, command, data = receive_response(
                    ser, max(0.1, deadline - time.monotonic())
                )
            except TimeoutError:
                break
            if command != CMD_FIRMWARE_ANNOUNCE_RESPONSE:
                continue
            if len(data) != 4:
                continue
            current_version, announced_version = struct.unpack(">HH", data)
            devices[source] = (current_version, announced_version)
            print(
                f"发现待升级设备：地址={source:04X}，"
                f"当前版本=0x{current_version:04X}，"
                f"目标版本=0x{announced_version:04X}"
            )

    print(f"共发现{len(devices)}台需要升级的设备。")
    print("人工确认后，请使用--target 设备短地址执行正式OTA。")


def main() -> int:
    """解析命令行参数并启动OTA测试。"""

    parser = argparse.ArgumentParser(description="STM8 Mesh继电器串口OTA测试工具")
    parser.add_argument("--port", required=True, help="电脑串口，例如COM3或/dev/cu.usbserial-xxx")
    parser.add_argument("--firmware", required=True, help="主程序裸BIN文件")
    parser.add_argument("--source", default="0001", type=lambda value: parse_hex_address(value, "source"),
                        help="主机Mesh短地址，默认0001")
    parser.add_argument("--target", default="FFFF", type=lambda value: parse_hex_address(value, "target"),
                        help="设备短地址，FFFF表示广播，默认FFFF")
    parser.add_argument("--baudrate", type=int, default=9600, help="UART波特率，默认9600")
    parser.add_argument("--timeout", type=float, default=3.0, help="普通应答超时秒数，默认3秒")
    parser.add_argument("--verify-timeout", type=float, default=30.0, help="CRC校验应答超时秒数，默认30秒")
    parser.add_argument("--result-timeout", type=float, default=20.0,
                        help="等待设备复位后上报升级结果的秒数，默认20秒")
    parser.add_argument("--discover", action="store_true",
                        help="只广播固件信息并收集待升级设备，不执行OTA")
    parser.add_argument("--discovery-timeout", type=float, default=3.0,
                        help="设备发现等待时间，默认3秒")
    parser.add_argument("--version-header", default="include/app_config.h",
                        help="版本配置头文件，默认include/app_config.h")
    parser.add_argument("--version", type=lambda value: int(value, 0), default=None,
                        help="可选，手工指定版本；不指定时自动读取FIRMWARE_VERSION")
    args = parser.parse_args()

    try:
        if args.discover:
            discover_upgrade_devices(args)
        else:
            send_ota(args)
    except (OSError, serial.SerialException, TimeoutError, ValueError, RuntimeError) as exc:
        print(f"OTA失败：{exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())