#!/usr/bin/env python3
"""
GD32E230 固件打包工具 v2 — 布局 profile 化双固件打包 (A/B 槽 + 记录/CRC)
================================================================

v2 模式(默认): 输入编译好的 APP bin, 自动识别 A/B 槽, 产出
  *_UPDATE.bin   A 槽 OTA/Ymodem 包 (整 N 页 + 末 4B trailer {页数16|CRC16})
  *_FULL.bin     整片镜像 (BL + A 包 + 槽尾锚点 + B 满槽 + 0xEEEE 标志)
  *_FACTORY.bin  工厂烧录镜像 (BL + B 满槽 + flag, 无 A 内容; 上电跳 B)
  *.md5.txt      产物 MD5 清单 + 布局摘要 (随主镜像生成)
打包规则与 BL (dev_fact v5.1) 逐字一致, 写盘后回读复核。计算与校验在
fw_pack.py 公共库 (FlashTool GUI 共用)。

legacy 模式(过渡期兼容, -f/-m 二选一时启用): 旧单槽方案
  -f/--fill  纯 0xFF 补页, 无 trailer (旧 BL 的 Ymodem 升级包)
  -m/--merge BL + APP + 0xEEEE 标志 拼整片镜像 (不补齐)

v2 Flash 布局(默认 profile CeramixNano_MB_Dual_FW, GD32E230F8/64K):
  0x08000000  Bootloader   8K   页 0~7
  0x08002000  FW_A 运行槽  36K  页 8~43   OTA 唯一接收区, 槽尾锚点@0x0800AFFC
  0x0800B000  FW_B 保底槽  19K  页 44~62  固定满槽, B 记录@0x0800FBF8
  0x0800FC00  标志页       1K   页 63     0xEEEE@0x0800FFFC
其他工程换 profile 即可(保底槽 8K/28K/无保底槽等), 见 --profile。
"""

import argparse
import os
import struct
import sys

import fw_pack

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, SCRIPT_DIR)

TOOL_VERSION = '2.1.0'

# 默认布局 profile (与 fw_pack.BUILTIN_PROFILES 键一致)
DEFAULT_PROFILE = 'CeramixNano_MB_Dual_FW'


# ═══════════════════════════════════════════════════════════
# v2 模式: 双固件 profile 打包
# ═══════════════════════════════════════════════════════════

def _collect_inputs(args, profile):
    """汇合 位置参数(自动识别) 与 --a/--b(显式指定), 冲突报错。
    返回 (a_path|None, b_path|None)。"""
    if (args.a or args.b) and args.bins:
        raise fw_pack.FwPackError(
            '位置输入(自动识别)与显式 --a/--b 不能混用, 二选一')
    if args.bins:
        found = {}
        for path in args.bins:
            if not os.path.isfile(path):
                raise fw_pack.FwPackError(f'找不到输入文件: {path}')
            data = open(path, 'rb').read()
            slot = fw_pack.detect_slot(data, profile)
            if slot in found:
                raise fw_pack.FwPackError(
                    f'{path!r} 与 {found[slot]!r} 都被识别为 {slot.upper()} 包, '
                    f'一个槽只能有一个输入')
            found[slot] = path
        return found.get('a'), found.get('b')
    if args.a and args.b and os.path.abspath(args.a) == os.path.abspath(args.b):
        raise fw_pack.FwPackError('--a 与 --b 指向同一个文件')
    for opt, path, slot in (('--a', args.a, 'a'), ('--b', args.b, 'b')):
        if not path:
            continue
        if not os.path.isfile(path):
            raise fw_pack.FwPackError(f'找不到 {opt} 输入文件: {path}')
        data = open(path, 'rb').read()
        which = fw_pack.detect_slot(data, profile)
        if which != slot:      # 显式指定与向量表判定冲突 → 报错(规格 4.1)
            raise fw_pack.FwPackError(
                f'{opt} 显式指定为 {slot.upper()} 槽, 但其复位向量落在 '
                f'{which.upper()} 槽 — 请改用 --{which} 或调整输入')
    return args.a, args.b


def run_v2(args):
    fw_pack.self_test()
    try:
        profile = fw_pack.load_profile(args.profile)
    except fw_pack.FwPackError as exc:
        print(f'错误：{exc}', file=sys.stderr)
        return 2
    try:
        a_path, b_path = _collect_inputs(args, profile)
        if not a_path and not b_path:
            print('错误：v2 模式至少需要一个输入 bin（A 包或 B 包）。\n'
                  '示例: python gd32e230_firmware_tool.py app_a.bin app_b.bin -b bl.bin',
                  file=sys.stderr)
            return 2
        out_dir = args.output
        if out_dir and os.path.splitext(args.output)[1].lower() == '.bin':
            print('错误：v2 模式 -o 是输出目录（产物自动命名）；'
                  'legacy 模式 -o 才是输出文件。', file=sys.stderr)
            return 2
        fw_pack.pack_all(
            args.profile, a_path, b_path=b_path, bl_path=args.bootloader,
            out_dir=out_dir, make_full=not args.no_full,
            make_factory=not args.no_factory)
    except fw_pack.FwPackError as exc:
        print(f'\n错误：{exc}', file=sys.stderr)
        return 1
    return 0


# ═══════════════════════════════════════════════════════════
# legacy 模式: 旧单槽方案 (纯 0xFF 补页 / 拼接, 无 trailer)
# ═══════════════════════════════════════════════════════════

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"
DEFAULT_BOOTLOADER = os.path.join(SCRIPT_DIR, "gd32e230f8_bootloader_hulk.bin")


def calc_layout(flash_size, app_addr=None):
    if app_addr is None:
        app_addr  = FLASH_BASE + BOOTLOADER_SIZE        # 0x08002000
        cal_addr  = FLASH_BASE + flash_size - CAL_PAGE_SIZE - FLAG_PAGE_SIZE
        app_max   = cal_addr - app_addr
        cal_reserved = True
    else:
        # 自定义落位（如 FW_B 槽 0x0800B000）：APP 可用到标志页前，校准页让渡给 APP
        cal_addr  = None
        app_max   = FLASH_BASE + flash_size - FLAG_PAGE_SIZE - app_addr
        cal_reserved = False
    flag_addr  = FLASH_BASE + flash_size - 4          # 最后一页末尾 - 4B
    flag_page  = FLASH_BASE + flash_size - FLAG_PAGE_SIZE
    return {
        "flash_size": flash_size,
        "flash_end":  FLASH_BASE + flash_size,
        "app_addr":   app_addr,
        "app_max":    app_max,
        "cal_addr":   cal_addr,
        "cal_reserved": cal_reserved,
        "flag_page":  flag_page,
        "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"]:
        limit_desc = (
            f"{layout['flash_size']//1024}KB - 8KB(BL) - 1KB(校准) - 1KB(Flag) "
            f"= {layout['app_max']//1024}KB"
            if layout["cal_reserved"] else
            f"APP 0x{layout['app_addr']:08X} ~ 0x{layout['flag_page']:08X}（Flag 页前）"
            f"= {layout['app_max']//1024}KB"
        )
        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: {limit_desc}",
            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}")
    if layout["cal_reserved"]:
        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)

    if len(padded) > layout["app_max"]:
        limit_desc = (
            f"{layout['flash_size']//1024}KB - 8KB(BL) - 1KB(校准) - 1KB(Flag) "
            f"= {layout['app_max']//1024}KB"
            if layout["cal_reserved"] else
            f"APP 0x{layout['app_addr']:08X} ~ 0x{layout['flag_page']:08X}（Flag 页前）"
            f"= {layout['app_max']//1024}KB"
        )
        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: {limit_desc}",
            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 run_legacy(args):
    chip = CHIP_CONFIG[args.flash_size]

    app_addr = None
    if args.app_addr is not None:
        try:
            app_addr = int(args.app_addr, 0)
        except ValueError:
            print(f"错误：-a/--app-addr 无法解析：{args.app_addr}", file=sys.stderr)
            return 2
        if app_addr % PAGE_SIZE != 0:
            print(f"错误：APP 落位地址 0x{app_addr:08X} 未按 1KB 页对齐", file=sys.stderr)
            return 2
        addr_min = FLASH_BASE + BOOTLOADER_SIZE
        addr_max = FLASH_BASE + chip["size"] - FLAG_PAGE_SIZE
        if not (addr_min <= app_addr < addr_max):
            print(
                f"错误：APP 落位地址 0x{app_addr:08X} 超出允许范围 "
                f"[0x{addr_min:08X}, 0x{addr_max:08X})",
                file=sys.stderr,
            )
            return 2

    layout = calc_layout(chip["size"], app_addr)
    print(f"[legacy] 芯片: {chip['label']}")

    app = args.bins[0] if args.bins else None
    if not app or not os.path.isfile(app):
        print(f"错误：找不到 APP 文件 {app}", file=sys.stderr)
        return 2

    if args.output is None:
        base, ext = os.path.splitext(os.path.basename(app))
        suffix = "_BL" if args.merge else "_UPDATE"
        args.output = os.path.join(os.path.dirname(os.path.abspath(app)),
                                   f"{base}{suffix}{ext}")

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


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

def main():
    parser = argparse.ArgumentParser(
        prog='gd32e230_firmware_tool',
        description=f"GD32E230 固件打包工具 v{TOOL_VERSION} — 双固件 profile 打包"
                    "（默认）/ 旧单槽补齐·拼接（-f/-m）",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=f"""
v2 双固件模式（默认, 规格: UPGRADE_PLAN_v2.md）:
  python %(prog)s app_a.bin                          # 识别 A 包 → *_UPDATE.bin (+FULL)
  python %(prog)s app_a.bin app_b.bin -b bl.bin      # 全产物: UPDATE/FULL/FACTORY/md5
  python %(prog)s app_a.bin --no-full                # 仅升级包
  python %(prog)s app_b.bin -b bl.bin                # 仅工厂镜像 → *_FACTORY.bin
  python %(prog)s --a a.bin --b b.bin                # 显式指定槽位(与向量表冲突报错)
  python %(prog)s a.bin --profile my_layout.json     # 自定义布局 profile JSON

legacy 模式（旧单槽方案, -f/-m 触发）:
  python %(prog)s -f app.bin                         # 纯 0xFF 补齐升级包(无 trailer)
  python %(prog)s -m app.bin                         # BL+APP+flag 拼整片
  python %(prog)s -m app.bin -s F6                   # F6 芯片
  python %(prog)s -m app.bin -a 0x0800B000 -b bl.bin # 自定义落位
        """,
    )
    parser.add_argument(
        "bins", nargs="*",
        help="v2: 输入 bin（A/B 包自动识别, 可多个）；legacy: APP 路径")
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument("-m", "--merge", action="store_true",
                      help="legacy: 合成完整 Flash 镜像 (bootloader + APP + FLAG)")
    mode.add_argument("-f", "--fill", action="store_true",
                      help="legacy: 仅补齐 APP 到整页边界 (纯 0xFF, 无 trailer)")
    parser.add_argument("--legacy", action="store_true",
                        help="显式声明 legacy 模式(需配合 -f/-m)")
    parser.add_argument("--profile", default=DEFAULT_PROFILE,
                        help="v2 布局 profile: 内置名(%s)或 JSON 文件路径"
                             % "/".join(fw_pack.BUILTIN_PROFILES))
    parser.add_argument("--a", metavar="BIN", default=None,
                        help="v2: 显式指定 A 包(运行槽固件)")
    parser.add_argument("--b", metavar="BIN", default=None,
                        help="v2: 显式指定 B 包(出厂保底固件)")
    parser.add_argument("-b", "--bootloader", default=None,
                        help="BL bin (v2: 参与 FULL 拼接; legacy -m: 默认内置 BL)")
    parser.add_argument("-o", "--output", default=None,
                        help="输出路径: v2=输出目录(自动命名); legacy=输出文件")
    parser.add_argument("--no-full", action="store_true",
                        help="v2: 不拼 FULL 整片镜像")
    parser.add_argument("--no-factory", action="store_true",
                        help="v2: 不出 FACTORY 工厂镜像")
    parser.add_argument("-s", "--flash-size", default=DEFAULT_CHIP,
                        choices=list(CHIP_CONFIG.keys()),
                        help="legacy 芯片型号: F4(16K) F6(32K) F8(64K)")
    parser.add_argument("-a", "--app-addr", default=None,
                        help="legacy: APP 落位地址, 如 0x0800B000 (默认 0x08002000)")

    args = parser.parse_args()
    print(f"gd32e230_firmware_tool v{TOOL_VERSION} (fw_pack {fw_pack.TOOL_VERSION})")

    legacy = args.fill or args.merge
    if args.legacy and not legacy:
        parser.error("--legacy 需配合 -f/--fill 或 -m/--merge 使用")
    if legacy and (args.profile != DEFAULT_PROFILE or args.a or args.b
                   or args.no_full or args.no_factory or args.app_addr):
        parser.error("--profile/--a/--b/--no-full/--no-factory/-a 是 v2/legacy "
                     "专属参数, 不能与 -f/-m 同用")
    if legacy and len(args.bins) != 1:
        parser.error("legacy 模式需要一个 APP 位置参数, 例: python %(prog)s -f app.bin"
                     % {'prog': parser.prog})

    rc = run_legacy(args) if legacy else run_v2(args)
    sys.exit(rc)


if __name__ == "__main__":
    main()