#!/usr/bin/env python3
"""
Generate a C++ cross-check header from the reference implementation (motalib).

Emits test/test_ota/mota_vectors.h containing a real .mota built by the Python tool plus the
expected parse results and a merkle proof, so the device-side C++ core (src/helpers/ota/) is
verified to agree byte-for-byte with the host packager.

Run:  ./meshcore/bin/python tools/mota/gen_vectors.py
"""
from __future__ import annotations

import random
from pathlib import Path

import motalib as ml

OUT = Path(__file__).resolve().parents[2] / "test" / "test_ota" / "mota_vectors.h"


def _carr(name, data: bytes) -> str:
    body = ",".join(str(b) for b in data)
    return f"static const uint8_t {name}[{len(data)}] = {{{body}}};\n"


def build_full():
    random.seed(42)
    fw = bytes(random.getrandbits(8) for _ in range(5 * 1024 + 137))  # 6 blocks, last short
    image, _ = ml.ensure_endf(fw)
    m = ml.build_manifest(
        target_id=0x11223344, fw_version=ml.pack_version("1.16.0"),
        image_size=len(image), payload=image, block_size=1024,
        image_hash=ml.mh32(image), codec_id=ml.CODEC_FULL, is_full=True, hw_id="TESTHW")
    return ml.build_container(m, image), m, image


def emit_proof_case(idx, count) -> str:
    """A tree of `count` arbitrary 4-byte leaves + every block's proof, computed by the reference."""
    random.seed(1000 + count)
    leaves = [bytes(random.getrandbits(8) for _ in range(4)) for _ in range(count)]
    root = ml.merkle_root(leaves)
    blob = b""
    offs, nsibs = [], []
    for i in range(count):
        sib = ml.proof_siblings(leaves, i)
        assert ml.verify_proof(leaves[i], i, ml.merkle_proof(leaves, i), root, count), (count, i)
        offs.append(len(blob))
        nsibs.append(len(sib) // 4)
        blob += sib
    p = f"T{idx}"
    out = [
        f"static const uint32_t {p}_COUNT = {count}u;",
        _carr(f"{p}_LEAVES", b"".join(leaves)),
        _carr(f"{p}_ROOT", root),
        f"static const uint16_t {p}_POFF[{count}] = {{{','.join(str(o) for o in offs)}}};",
        f"static const uint8_t  {p}_PNSIB[{count}] = {{{','.join(str(n) for n in nsibs)}}};",
        _carr(f"{p}_PBLOB", blob),
    ]
    return "\n".join(out)


def main():
    blob, m, image = build_full()
    leaves = ml.leaf_hashes(image, 1024)
    assert ml.merkle_root(leaves) == m.merkle_root

    proof_idx = 2 if m.block_count > 2 else 0
    siblings = ml.proof_siblings(leaves, proof_idx)
    assert ml.verify_proof(leaves[proof_idx], proof_idx, ml.merkle_proof(leaves, proof_idx),
                           m.merkle_root, len(leaves))

    lines = [
        "// AUTO-GENERATED by tools/mota/gen_vectors.py — do not edit by hand.",
        "// Cross-check vectors: a real .mota from the reference packager (motalib.py).",
        "#pragma once",
        "#include <stdint.h>",
        "",
        _carr("MOTA_VEC", blob),
        f"static const uint32_t MOTA_VEC_LEN = {len(blob)};",
        f"static const uint32_t EXP_TARGET_ID = 0x{m.target_id:08x}u;",
        f"static const uint32_t EXP_FW_VERSION = 0x{m.fw_version:08x}u;",
        f"static const uint32_t EXP_IMAGE_SIZE = {m.image_size}u;",
        f"static const uint32_t EXP_PAYLOAD_SIZE = {m.payload_size}u;",
        f"static const uint32_t EXP_BLOCK_COUNT = {m.block_count}u;",
        f"static const uint8_t  EXP_BLOCK_SIZE_LOG2 = {m.block_size_log2};",
        f"static const uint8_t  EXP_CODEC_ID = {m.codec_id};",
        _carr("EXP_MERKLE_ROOT", m.merkle_root),
        _carr("EXP_IMAGE_HASH", m.image_hash),
        _carr("EXP_HW_ID", m.hw_id),
        f"static const uint32_t PROOF_INDEX = {proof_idx}u;",
        f"static const uint8_t  PROOF_NSIB = {len(siblings)//4};",
        _carr("PROOF_SIBLINGS", siblings),
    ]

    # exhaustive per-index proof cases for several tricky counts (deep promotion chains)
    counts = [5, 7, 8, 65, 100, 255, 256]
    for i, c in enumerate(counts):
        lines.append(emit_proof_case(i, c))
    lines.append("struct ProofCase { uint32_t count; const uint8_t* leaves; const uint8_t* root;"
                 " const uint16_t* poff; const uint8_t* pnsib; const uint8_t* pblob; };")
    rows = ",".join(f"{{T{i}_COUNT,T{i}_LEAVES,T{i}_ROOT,T{i}_POFF,T{i}_PNSIB,T{i}_PBLOB}}"
                    for i in range(len(counts)))
    lines.append(f"static const ProofCase PROOF_CASES[] = {{{rows}}};")
    lines.append(f"static const int N_PROOF_CASES = {len(counts)};")

    # small-block signed .mota for the host transfer simulation (each block fits one packet => no
    # fragmentation needed to validate the manager/protocol end-to-end). Deterministic signing key.
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
    random.seed(99)
    sim_fw = bytes(random.getrandbits(8) for _ in range(2000))
    sim_image, _ = ml.ensure_endf(sim_fw)
    sim_priv = Ed25519PrivateKey.from_private_bytes(bytes(range(32)))
    sm = ml.build_manifest(target_id=0xCAFEBABE, fw_version=ml.pack_version("3.0.0"),
                           image_size=len(sim_image), payload=sim_image, block_size=128,
                           image_hash=ml.mh32(sim_image), codec_id=ml.CODEC_FULL,
                           is_full=True, sign_priv=sim_priv)
    sim_blob = ml.build_container(sm, sim_image)
    lines.append("// small-block signed .mota for the OtaManager host transfer simulation")
    lines.append(_carr("SIM_MOTA", sim_blob))
    lines.append(f"static const uint32_t SIM_MOTA_LEN = {len(sim_blob)};")
    lines.append(f"static const uint32_t SIM_TARGET_ID = 0x{sm.target_id:08x}u;")

    # 1 KB-block signed .mota: each logical block spans MULTIPLE LoRa DATA fragments (1024 / 160 = 7),
    # so this exercises the multi-fragment reassembly + split proof path the device actually uses.
    sm1k = ml.build_manifest(target_id=0xCAFEBABE, fw_version=ml.pack_version("3.0.0"),
                             image_size=len(sim_image), payload=sim_image, block_size=1024,
                             image_hash=ml.mh32(sim_image), codec_id=ml.CODEC_FULL,
                             is_full=True, sign_priv=sim_priv)
    sim1k_blob = ml.build_container(sm1k, sim_image)
    lines.append("// 1 KB-block signed .mota (multi-fragment per block) for the reassembly transfer test")
    lines.append(_carr("SIM_MOTA_1K", sim1k_blob))
    lines.append(f"static const uint32_t SIM_MOTA_1K_LEN = {len(sim1k_blob)};")
    lines.append(f"static const uint32_t SIM_MOTA_1K_BLOCKS = {sm1k.block_count};")

    # detools sequential+crle delta vector: base image, the real detools 0.53.0 patch, and the
    # expected target image. The native test applies DT_PATCH to DT_BASE with the *vendored detools
    # C decoder* (src/helpers/ota/detools) and must reproduce DT_TARGET byte-for-byte -- proving the
    # on-device delta apply path uses detools, not a reimplementation.
    import io
    import detools
    random.seed(2024)
    dt_base_body = bytes(random.getrandbits(8) for _ in range(3000))
    dt_tgt = bytearray(dt_base_body)
    for off in (137, 138, 139, 1500, 2999):          # localized edits, version-bump style
        dt_tgt[off] ^= 0x5A
    dt_tgt += bytes(random.getrandbits(8) for _ in range(200))   # small appended tail
    dt_base_img, _ = ml.ensure_endf(dt_base_body)
    dt_tgt_img, _ = ml.ensure_endf(bytes(dt_tgt))
    fp = io.BytesIO()
    detools.create_patch(io.BytesIO(dt_base_img), io.BytesIO(dt_tgt_img), fp,
                         patch_type="sequential", compression="crle")
    dt_patch = fp.getvalue()
    lines.append("// detools sequential+crle delta: apply DT_PATCH to DT_BASE -> DT_TARGET")
    lines.append(_carr("DT_BASE", dt_base_img))
    lines.append(f"static const uint32_t DT_BASE_LEN = {len(dt_base_img)};")
    lines.append(_carr("DT_PATCH", dt_patch))
    lines.append(f"static const uint32_t DT_PATCH_LEN = {len(dt_patch)};")
    lines.append(_carr("DT_TARGET", dt_tgt_img))
    lines.append(f"static const uint32_t DT_TARGET_LEN = {len(dt_tgt_img)};")

    # detools IN-PLACE+crle delta (nRF52 single-slot apply): apply DT_IP_PATCH to a memory region
    # holding DT_IP_BASE -> region[0:to_size] == DT_IP_TARGET. The native test runs the vendored
    # in-place decoder over a DT_IP_MEM-byte RAM buffer, exactly as the bootloader will over flash.
    DT_IP_MEM, DT_IP_SEG = 0x8000, 0x1000
    ip_patch_io = io.BytesIO()
    detools.create_patch(io.BytesIO(dt_base_img), io.BytesIO(dt_tgt_img), ip_patch_io,
                         patch_type="in-place", memory_size=DT_IP_MEM, segment_size=DT_IP_SEG,
                         compression="crle")
    dt_ip_patch = ip_patch_io.getvalue()
    lines.append("// detools in-place+crle delta: apply DT_IP_PATCH over a DT_IP_MEM buffer holding DT_IP_BASE")
    lines.append(_carr("DT_IP_BASE", dt_base_img))
    lines.append(f"static const uint32_t DT_IP_BASE_LEN = {len(dt_base_img)};")
    lines.append(_carr("DT_IP_PATCH", dt_ip_patch))
    lines.append(f"static const uint32_t DT_IP_PATCH_LEN = {len(dt_ip_patch)};")
    lines.append(_carr("DT_IP_TARGET", dt_tgt_img))
    lines.append(f"static const uint32_t DT_IP_TARGET_LEN = {len(dt_tgt_img)};")
    lines.append(f"static const uint32_t DT_IP_MEM = {DT_IP_MEM}u;")

    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text("\n".join(lines) + "\n")
    print(f"wrote {OUT}")
    print(f"  blob={len(blob)}B  blocks={m.block_count}  proof_idx={proof_idx} nsib={len(siblings)//4}")
    print(f"  proof cases: counts={counts}")
    print(f"  detools delta: base={len(dt_base_img)}B target={len(dt_tgt_img)}B seq_patch={len(dt_patch)}B inplace_patch={len(dt_ip_patch)}B")


if __name__ == "__main__":
    main()
