#!/usr/bin/env python3
"""
GD32E230 固件工具：整页补齐 + 合成烧录镜像，二合一。

模式（互斥）：
  -m / --merge   合成完整 Flash 镜像 (bootloader + APP + FLAG)
                 → 用于 OpenOCD / J-Link 整片烧录，不补齐
  -f / --fill    仅补齐 APP 到整页边界
                 → 用于 Bootloader Ymodem 升级

Flash 布局：
  0x08000000  ┌──────────────┐
              │  Bootloader  │  8KB
  APP_ADDR    ├──────────────┤
              │  Application │  最大 54KB (F8) / 22KB (F6) / 6KB (F4)
              │     ...      │
  CAL_ADDR    │  校准数据     │  1KB (预留)
  FLAG_ADDR   │  0xEEEE 标志 │  1KB
  FLASH_END   └──────────────┘
"""

import argparse
import struct
import sys
import os


# ═══════════════════════════════════════════════════════════
# 常量
# ═══════════════════════════════════════════════════════════

FLASH_BASE       = 0x08000000
BOOTLOADER_SIZE  = 8 * 1024                     # 8KB
CAL_PAGE_SIZE    = 1 * 1024                     # 1KB 校准数据预留
FLAG_PAGE_SIZE   = 1 * 1024                     # 1KB Flag 页
PAGE_SIZE        = 1024                         # Flash 页大小 (1KB)
FILL_BYTE        = 0xFF                         # 补齐字节
FLAG_VALUE       = 0xEEEE                       # 与 bootloader 统一

CHIP_CONFIG = {
    "F4": {"size": 16 * 1024, "label": "GD32E230F4 (16KB)"},
    "F6": {"size": 32 * 1024, "label": "GD32E230F6 (32KB)"},
    "F8": {"size": 64 * 1024, "label": "GD32E230F8 (64KB)"},
}

DEFAULT_CHIP       = "F8"
SCRIPT_DIR         = os.path.dirname(os.path.abspath(__file__))
DEFAULT_BOOTLOADER = os.path.join(SCRIPT_DIR, "gd32e230f8_bootloader_hulk.bin")


# ═══════════════════════════════════════════════════════════
# 地址计算
# ═══════════════════════════════════════════════════════════

def calc_layout(flash_size):
    app_addr  = FLASH_BASE + BOOTLOADER_SIZE        # 0x08002000
    cal_addr  = FLASH_BASE + flash_size - CAL_PAGE_SIZE - FLAG_PAGE_SIZE
    flag_addr = FLASH_BASE + flash_size - 4          # 最后一页末尾 - 4B
    app_max   = cal_addr - app_addr
    return {
        "flash_size": flash_size,
        "flash_end":  FLASH_BASE + flash_size,
        "app_addr":   app_addr,
        "app_max":    app_max,
        "cal_addr":   cal_addr,
        "flag_addr":  flag_addr,
    }


# ═══════════════════════════════════════════════════════════
# 补齐
# ═══════════════════════════════════════════════════════════

def pad_to_page(data, page_size=PAGE_SIZE, fill_byte=FILL_BYTE):
    remainder = len(data) % page_size
    if remainder == 0:
        return data, 0
    pad_len = page_size - remainder
    return data + bytes([fill_byte]) * pad_len, pad_len


# ═══════════════════════════════════════════════════════════
# 合成模式
# ═══════════════════════════════════════════════════════════

def merge(bootloader_path, app_path, output_path, layout):
    # 读取
    with open(bootloader_path, "rb") as f:
        bl = f.read()
    with open(app_path, "rb") as f:
        app = f.read()

    # 校验
    app_offset = layout["app_addr"] - FLASH_BASE

    if len(bl) > app_offset:
        print(f"错误：Bootloader ({len(bl)}B) 超过 APP 偏移 ({app_offset}B)", file=sys.stderr)
        return False
    if len(app) > layout["app_max"]:
        print(
            f"错误：APP ({len(app)}B / {len(app)/1024:.1f}KB) 超过最大容量 "
            f"({layout['app_max']}B / {layout['app_max']/1024:.0f}KB)\n"
            f"  芯片 {layout['flash_size']//1024}KB: "
            f"{layout['flash_size']//1024}KB - 8KB(BL) - 1KB(校准) - 1KB(Flag) "
            f"= {layout['app_max']//1024}KB",
            file=sys.stderr,
        )
        return False

    # 构建镜像
    image = bytearray(b'\xFF' * layout["flash_size"])
    image[0:len(bl)] = bl
    image[app_offset:app_offset + len(app)] = app
    flag_offset = layout["flag_addr"] - FLASH_BASE
    struct.pack_into("<I", image, flag_offset, FLAG_VALUE)

    # 输出
    with open(output_path, "wb") as f:
        f.write(image)

    print(f"\n合成完成: {output_path}")
    print(f"  芯片       : {layout['flash_size']//1024}KB")
    print(f"  Bootloader : {len(bl):>6}B  @ 0x{FLASH_BASE:08X}")
    print(f"  Application: {len(app):>6}B  @ 0x{layout['app_addr']:08X}")
    print(f"  Flag       : 0x{FLAG_VALUE:04X}     @ 0x{layout['flag_addr']:08X}")
    print(f"  校准预留   : 1KB      @ 0x{layout['cal_addr']:08X}")
    print(f"  APP 剩余   : {layout['app_max'] - len(app):>6}B")
    print(f"  总大小     : {layout['flash_size']:>6}B  ({layout['flash_size']//1024}KB)")
    return True


# ═══════════════════════════════════════════════════════════
# 补齐模式
# ═══════════════════════════════════════════════════════════

def fill_only(input_path, output_path, layout):
    with open(input_path, "rb") as f:
        data = f.read()

    padded, pad_len = pad_to_page(data)

    # 校验补齐后不超过 APP 最大容量
    if len(padded) > layout["app_max"]:
        print(
            f"错误：补齐后 ({len(padded)}B / {len(padded)/1024:.1f}KB) 超过 APP 最大容量 "
            f"({layout['app_max']}B / {layout['app_max']/1024:.0f}KB)\n"
            f"  芯片 {layout['flash_size']//1024}KB: "
            f"{layout['flash_size']//1024}KB - 8KB(BL) - 1KB(校准) - 1KB(Flag) "
            f"= {layout['app_max']//1024}KB",
            file=sys.stderr,
        )
        sys.exit(1)

    with open(output_path, "wb") as f:
        f.write(padded)

    print(f"补齐完成: {output_path}")
    print(f"  原始大小: {len(data)}B")
    print(f"  页大小  : {PAGE_SIZE}B (1KB)")
    print(f"  补齐字节: {pad_len}B")
    print(f"  输出大小: {len(padded)}B")
    print(f"  APP 上限: {layout['app_max']}B ({layout['app_max']//1024}KB) → "
          f"剩余 {layout['app_max'] - len(padded)}B")

    if pad_len == 0:
        print(f"  说明    : 已是整页对齐，原样输出")


# ═══════════════════════════════════════════════════════════
# 入口
# ═══════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(
        description="GD32E230 固件工具 — 整页补齐 & 合成烧录镜像",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python %(prog)s -m app.bin                     # F8 合成镜像 (默认)
  python %(prog)s -m app.bin -s F6               # F6 合成镜像
  python %(prog)s -f app.bin                     # F8 补齐升级包 (默认)
  python %(prog)s -f app.bin -s F4               # F4 补齐升级包
  python %(prog)s -m app.bin -b my_bl.bin -o out.bin  # 自定义 bootloader/输出
        """,
    )

    # 位置参数
    parser.add_argument(
        "app", help="Application 的 .bin 文件路径"
    )

    # 模式 (互斥)
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument(
        "-m", "--merge",
        action="store_true",
        help="合成完整 Flash 镜像 (bootloader + APP + FLAG)，用于整片烧录",
    )
    mode.add_argument(
        "-f", "--fill",
        action="store_true",
        help="仅补齐 APP 到整页边界，用于 Bootloader Ymodem 升级",
    )

    # 选项
    parser.add_argument(
        "-s", "--flash-size",
        default=DEFAULT_CHIP,
        choices=list(CHIP_CONFIG.keys()),
        help=f"芯片型号: F4(16K) F6(32K) F8(64K) (默认: {DEFAULT_CHIP})",
    )
    parser.add_argument(
        "-b", "--bootloader",
        default=DEFAULT_BOOTLOADER,
        help=f"Bootloader bin 路径 (默认: {DEFAULT_BOOTLOADER})",
    )
    parser.add_argument(
        "-o", "--output",
        default=None,
        help="输出文件路径 (默认: APP名_BL.bin 或 APP名_UPDATE.bin)",
    )

    args = parser.parse_args()

    chip = CHIP_CONFIG[args.flash_size]
    layout = calc_layout(chip["size"])

    print(f"芯片: {chip['label']}")

    # 输出文件名
    if args.output is None:
        base, ext = os.path.splitext(os.path.basename(args.app))
        suffix = "_BL" if args.merge else "_UPDATE"
        args.output = f"{base}{suffix}{ext}"

    # 输入文件检查
    if not os.path.isfile(args.app):
        print(f"错误：找不到 APP 文件 {args.app}", file=sys.stderr)
        sys.exit(1)

    if args.merge:
        if not os.path.isfile(args.bootloader):
            print(f"错误：找不到 Bootloader 文件 {args.bootloader}", file=sys.stderr)
            sys.exit(1)
        if not merge(args.bootloader, args.app, args.output, layout):
            sys.exit(1)
    else:
        fill_only(args.app, args.output, layout)


if __name__ == "__main__":
    main()