"""Rebuild the replacement payload for the `os.execute` + `dd` route,
in a location-independent layout.

Layout produced:
  <out>/payload/            NNN_<name>.bin  +  payload.lua   (bare filenames)
  <out>/payload_rollback/   NNN_<name>.bin  +  payload.lua
  <out>/flash.sh            linear nsh script (no $( ), no && , no globbing)
  <out>/rollback.sh
  <out>/payload.json        machine-readable manifest

All offsets are multiples of 512 and every emitted length is a multiple of 4,
so `dd bs=4 seek=offset/4 count=length/4` writes exact ranges with no partial
trailing block.
"""
import os, sys, json, hashlib
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from paths import PROJ, INPUTS, WORK
STOCK = os.path.join(INPUTS, 'target', 'vela_app.bin')
PATCHED = os.path.join(WORK, 's4_from_s5', 'vela_app_s5icons.bin')
OUT = PROJ

recs_in = json.load(open(os.path.join(WORK, 's4_from_s5', 'manifest.json'), encoding='utf-8'))
stock = open(STOCK, 'rb').read()
patched = open(PATCHED, 'rb').read()
assert len(stock) == len(patched)

records = [dict(name='canary_superblock', offset=0, length=32,
                data=stock[:32], changed=True)]
assert stock[:8] == b'-rom1fs-'
for r in recs_in['records']:
    off, ln = r['offset'], r['length']
    records.append(dict(name=r['name'], offset=off, length=ln,
                        data=patched[off:off + ln], changed=stock[off:off + ln] != patched[off:off + ln]))

emitted = [r for r in records if r['changed']]
dropped = [r for r in records if not r['changed']]


def write_pack(dirname, source):
    d = os.path.join(OUT, dirname)
    os.makedirs(d, exist_ok=True)
    for f in os.listdir(d):
        os.remove(os.path.join(d, f))
    rows = []
    for i, r in enumerate(emitted):
        fn = f'{i:03d}_{r["name"]}.bin'
        open(os.path.join(d, fn), 'wb').write(source(r))
        assert r['offset'] % 4 == 0, r['name']          # 偏移恒定 4 对齐（实际是 512）
        rows.append(dict(name=r['name'], file=fn, offset=r['offset'], length=r['length']))
    with open(os.path.join(d, 'payload.lua'), 'w', encoding='utf-8', newline='\n') as f:
        f.write('-- generated by _tools/make_dd_payload.py -- do not hand-edit\n')
        f.write('-- dd: -if <this dir>/<file> -of /dev/app bs=4 seek=offset/4 count=length/4\n')
        f.write('return {\n')
        for r in rows:
            f.write(f'  {{ name = "{r["name"]}", file = "{r["file"]}", '
                    f'offset = {r["offset"]}, length = {r["length"]} }},\n')
        f.write('}\n')
    return rows


lua_rows = write_pack('payload', lambda r: r['data'])
r_rows = write_pack('payload_rollback',
                    lambda r: stock[r['offset']:r['offset'] + r['length']])

# ---------------------------------------------------------------- 合并成「单个条目」
# 关键：.face 里**每个文件 = 一个资源条目**。逐条打包会做出 480 个条目，
# 把表盘条目数顶到宿主吃不下的量（实测 6.59MB/170 条能跑，9.12MB/490 条黑屏；
# 原厂 resource.bin 到 9.43MB 都能跑 -> 卡的是条目数，不是体积）。
# 所以给编译用的是一份「魔数 + 所有负载顺序拼接」的单文件，偏移在构建时按相对偏移回填。
MAGIC_PAY = b'\x00S4E-PAYv1\x00\xa5\x5a\xc3\x3c'      # 15 B，必须全局唯一
MAGIC_RB = b'\x00S4E-RBKv1\x00\x5a\xa5\x3c\xc3'
BLOB = os.path.join(OUT, 'blob')
os.makedirs(BLOB, exist_ok=True)
for f in os.listdir(BLOB):
    os.remove(os.path.join(BLOB, f))

pay_blob = MAGIC_PAY + b''.join(r['data'] for r in emitted)
rb_blob = MAGIC_RB + b''.join(stock[r['offset']:r['offset'] + r['length']] for r in emitted)
open(os.path.join(BLOB, 'payload.bin'), 'wb').write(pay_blob)
open(os.path.join(BLOB, 'rollback.bin'), 'wb').write(rb_blob)

rel, off = [], 0
for r in emitted:
    rel.append(dict(name=r['name'], rel=off, target=r['offset'], length=r['length']))
    off += r['length']
json.dump(dict(magic_len=len(MAGIC_PAY), magic_rb_len=len(MAGIC_RB),
               records=rel, pay_bytes=len(pay_blob), rb_bytes=len(rb_blob)),
          open(os.path.join(BLOB, 'offsets.json'), 'w', encoding='utf-8'), indent=1)
print('blob/ 单条目打包: payload.bin %d B + rollback.bin %d B（%d 条记录，相对偏移表已写）'
      % (len(pay_blob), len(rb_blob), len(rel)))

json.dump(dict(device='Xiaomi Watch S4 eSIM (o62lte)', target='/dev/app',
               dd="dd if=<dir>/<file> of=/dev/app bs=4 seek=<offset/4> count=<length/4>",
               block_size=4, records=lua_rows,
               dropped_identical=[r['name'] for r in dropped]),
          open(os.path.join(OUT, 'payload.json'), 'w'), indent=1, ensure_ascii=False)


def script(fname, rows, header, base, verify, tail):
    with open(os.path.join(OUT, fname), 'w', encoding='utf-8', newline='\n') as f:
        f.write('#!/bin/sh\n')
        for line in header:
            f.write('# ' + line + '\n')
        f.write('\nDEV=/dev/app\nBASE=' + base + '\n\n')
        f.write('echo "target check"\nls -l $DEV\n\n')
        for r in rows:
            # RLE 素材的槽位长度常常不是 4 的倍数 -> 那条退化成 bs=1（偏移本身就对齐）
            b = 4 if (r['offset'] % 4 == 0 and r['length'] % 4 == 0) else 1
            f.write('echo "== ' + r['name'] + ' ==\n')
            f.write(f'dd if=$BASE/{r["file"]} of=$DEV bs={b} seek={r["offset"]//b} '
                    f'count={r["length"]//b}\n')
            if verify:
                f.write(f'dd if=$DEV of=/tmp/chk.bin bs={b} skip={r["offset"]//b} '
                        f'count={r["length"]//b}\n')
                f.write(f'cmp /tmp/chk.bin $BASE/{r["file"]}\n\n')
        f.write(tail + '\n')


script('flash.sh', lua_rows,
       ['S5 -> S4 icon replacement, run on-device in nsh (NuttX shell).',
        '把 flash.sh 与 payload/ 放到同一个目录，例如 /data/s5icons/',
        '然后:  sh /data/s5icons/flash.sh     最后:  reboot',
        'nsh 不支持 $( ) / && / || / 通配符，所以这里是线性展开的。'],
       '/data/s5icons/payload', True,
       'rm /tmp/chk.bin\necho "all done - run reboot to apply"')

script('rollback.sh', r_rows,
       ['写回 S4 原厂图标'], '/data/s5icons/payload_rollback', False,
       'echo "rollback done - run reboot to apply"')

total = sum(r['length'] for r in lua_rows)
print(f'emitted {len(lua_rows)} records, {total} bytes ({total/1048576:.2f} MB)')
print(f'dropped {len(dropped)} byte-identical record(s): {[r["name"] for r in dropped]}')
for r in lua_rows[:3]:
    print(f'  {r["name"]:<22} off=0x{r["offset"]:08x} len={r["length"]:<7} seek={r["offset"]//4} count={r["length"]//4}')
print('  ...')