# -*- coding: utf-8 -*-
"""按「角色」切负载：replace（只负责替换）/ restore（只负责还原）。

为什么拆成两张盘
----------------
「负载 + 回滚」合在一张盘 = 两份 15.5 MB ≈ 31 MB，超过 .face 的体积上限。
拆开之后每张盘只装一份负载（≈15.5 MB）。

拆盘后版本校验怎么办
--------------------
每张盘再带一张**指纹表**：每条记录存「另一侧字节的前 64 B」。
  * replace 盘：pay = q63 的帧（新字节）；fp = **原厂字节**前 64 B
  * restore 盘：pay = 原厂字节；      fp = **新字节**前 64 B
于是「写入前版本校验」和「写入后回读比对」两道保险都还在，代价只有 119*64 ≈ 7.4 KB。

关键约定（S4e 那版踩过的坑）
----------------------------
Lua 表里存的是**分区键名**（'health'），运行期才由 `DEVS[键]` 拼成 `/dev/health`。
若把 '/dev/health' 直接存进表里，拼出来就是 `if=nil` 或 `if=/dev//dev/health`。

产出
----
  payload/          NNN_<name>.bin + payload.lua          （replace 盘的负载）
  payload_restore/  NNN_<name>.bin + payload.lua          （restore 盘的负载）
  blob_replace/     payload.bin（魔数 + 顺序拼接）+ fp.bin + offsets.json
  blob_restore/     同上
  payload.json      机器可读清单
"""
import os, sys, json, hashlib, struct
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from paths import PROJ, WORK, TARGET_IMAGE

ROLE = (sys.argv[1] if len(sys.argv) > 1 else 'replace').lower()
assert ROLE in ('replace', 'restore'), ROLE
PATCHED = os.path.join(WORK, 'anim', 'vela_health_anim.bin')
OUT = PROJ
DEV = 'health'                     # ← Lua 表里的键名（不是 '/dev/health'）
FP_LEN = 64

recs_in = json.load(open(os.path.join(WORK, 'anim', 'manifest.json'), encoding='utf-8'))
stock = open(TARGET_IMAGE, 'rb').read()
patched = open(PATCHED, 'rb').read()
assert len(stock) == len(patched), '改后镜像长度变了 —— 等长覆盖的前提被破坏'
assert stock[:8] == b'-rom1fs-'

SRC = patched if ROLE == 'replace' else stock       # 要写进设备的字节
OTHER = stock if ROLE == 'replace' else patched     # 另一侧字节（做指纹）
FPDIR = os.path.join(OUT, 'payload' if ROLE == 'replace' else 'payload_restore')
BLOBDIR = os.path.join(OUT, 'blob_' + ROLE)

records = []
for r in recs_in['records']:
    off, ln = r['offset'], r['length']
    records.append(dict(name=r['name'], offset=off, length=ln, data=SRC[off:off + ln],
                        changed=(stock[off:off + ln] != patched[off:off + ln]),
                        fp=OTHER[off:off + FP_LEN]))
# replace 盘跳过「本来就没变」的记录（省体积）；restore 盘一条都不能少
emitted = [r for r in records if (r['changed'] or ROLE == 'restore')]
dropped = [r for r in records if r not in emitted]
for r in emitted:
    assert len(r['data']) == r['length'], r['name']
    assert len(r['fp']) == FP_LEN, r['name']


def write_pack(d):
    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 = '%03d_%s.bin' % (i, r['name'])
        open(os.path.join(d, fn), 'wb').write(r['data'])
        assert r['offset'] % 4 == 0, r['name']
        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/s5e_payload.py (%s) -- do not hand-edit\n' % ROLE)
        f.write('-- dd: -if <this dir>/<file> -of /dev/%s bs=4 seek=offset/4 count=length/4\n' % DEV)
        f.write('return {\n')
        for r in rows:
            f.write('  { name = "%s", file = "%s", offset = %d, length = %d },\n'
                    % (r['name'], r['file'], r['offset'], r['length']))
        f.write('}\n')
    return rows


lua_rows = write_pack(FPDIR)

# 魔数必须全局唯一，且**长度不要手写常量**（S4e/S5e 图标包都因为手写成 16 导致偏移差 1）
MAGIC_PAY = b'\x00S5P-PAYv1\x00\xa5\x5a\xc3\x3c'
MAGIC_FP = b'\x00S5P-FPv1\x00\x5a\xa5\x3c\xc3'
os.makedirs(BLOBDIR, exist_ok=True)
for f in os.listdir(BLOBDIR):
    os.remove(os.path.join(BLOBDIR, f))
pay_blob = MAGIC_PAY + b''.join(r['data'] for r in emitted)
fp_blob = MAGIC_FP + b''.join(r['fp'] for r in emitted)
open(os.path.join(BLOBDIR, 'payload.bin'), 'wb').write(pay_blob)
open(os.path.join(BLOBDIR, 'fp.bin'), 'wb').write(fp_blob)
rel, off, fpoff = [], 0, 0
for r in emitted:
    # 每条都存独立 sha256：.face 构建时用它校验"确实原样连续存放"，
    # 而不是拿 .face 跟自己比（那种检查恒真、抓不到 bug）
    rel.append(dict(name=r['name'], rel=off, fprel=fpoff, target=r['offset'],
                    length=r['length'], fplen=len(r['fp']),
                    pay_sha=hashlib.sha256(pay_blob[len(MAGIC_PAY) + off:
                                                  len(MAGIC_PAY) + off + r['length']]).hexdigest(),
                    fp_sha=hashlib.sha256(fp_blob[len(MAGIC_FP) + fpoff:
                                                 len(MAGIC_FP) + fpoff + len(r['fp'])]).hexdigest()))
    off += r['length']
    fpoff += len(r['fp'])
assert fpoff == len(fp_blob) - len(MAGIC_FP), (fpoff, len(fp_blob))
assert off == len(pay_blob) - len(MAGIC_PAY), (off, len(pay_blob))
json.dump(dict(role=ROLE, device=DEV, magic_len=len(MAGIC_PAY), magic_fp_len=len(MAGIC_FP),
               fp_len=FP_LEN, records=rel, pay_bytes=len(pay_blob), fp_bytes=len(fp_blob),
               dropped_identical=[r['name'] for r in dropped]),
          open(os.path.join(BLOBDIR, 'offsets.json'), 'w', encoding='utf-8'), indent=1)
print('blob_%s/ 单条目: payload.bin %d B + fp.bin %d B（%d 条记录，跳过 %d 条未变）'
      % (ROLE, len(pay_blob), len(fp_blob), len(rel), len(dropped)))

json.dump(dict(device='Xiaomi Watch S5 eSIM 46mm (p62lte)', target='/dev/' + DEV, dev_key=DEV,
               partition='vela_health.bin', role=ROLE, kind='pressure-measure-anim',
               dd='dd if=<dir>/<file> of=/dev/%s bs=4 seek=<offset/4> count=<length/4>' % DEV,
               block_size=4, records=lua_rows,
               dropped_identical=[r['name'] for r in dropped]),
          open(os.path.join(OUT, 'payload.json'), 'w', encoding='utf-8'), indent=1,
          ensure_ascii=False)
total = sum(r['length'] for r in lua_rows)
print('role=%s: emitted %d records, %.2f MB; dropped %d'
      % (ROLE, len(lua_rows), total / 1048576, len(dropped)))