#!/usr/bin/env python3
"""Correct the navigation display-speed source once, before it is shared."""

from __future__ import annotations

import argparse
import hashlib
import struct
from pathlib import Path


SOURCE_SHA256 = "f284d61e18f58f10db032e64799f7c3c348d8533c5b48e2ef390b69606a81648"
SETTER_OFFSET = 0x018B302C
SETTER_CONTINUE = SETTER_OFFSET + 4
SETTER_ORIGINAL = bytes.fromhex("00 f5 ad 30")
CAVE_OFFSET = 0x0189DAB4
CAVE_SIZE = 60
DISPLAY_CALL_SITES = (
    0x019F5376,
    0x019F9F70,
    0x019FA034,
    0x019FA0F8,
    0x01A23E6E,
)
CONVERT_SPEED = 0x018EC9E8
SAFETY_CALL_SITE = 0x020C411A


def digest(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def encode_thumb_branch(source: int, destination: int, link: bool = False) -> bytes:
    displacement = destination - (source + 4)
    if displacement & 1:
        raise ValueError("Thumb branch destination is not halfword aligned")
    if not -(1 << 24) <= displacement < (1 << 24):
        raise ValueError("Thumb branch destination is out of range")
    immediate = displacement & 0x01FFFFFF
    sign = (immediate >> 24) & 1
    i1 = (immediate >> 23) & 1
    i2 = (immediate >> 22) & 1
    j1 = (~(i1 ^ sign)) & 1
    j2 = (~(i2 ^ sign)) & 1
    first = 0xF000 | (sign << 10) | ((immediate >> 12) & 0x03FF)
    second_base = 0xD000 if link else 0x9000
    second = second_base | (j1 << 13) | (j2 << 11) | ((immediate >> 1) & 0x07FF)
    return struct.pack("<HH", first, second)


def source_wrapper() -> bytes:
    # r1 contains the incoming IEEE-754 display-speed float. Keep values below
    # 1.0 unchanged; otherwise subtract exactly 1.0, then execute the displaced
    # stock object-offset instruction and return to the unmodified setter body.
    body = bytes.fromhex(
        "07 ee 90 1a"  # vmov s15, r1
        " b7 ee 00 7a"  # vmov.f32 s14, #1.0
        " f4 ee c7 7a"  # vcmpe.f32 s15, s14
        " f1 ee 10 fa"  # vmrs APSR_nzcv, fpscr
        " 03 db"  # blt skip_subtract
        " 77 ee c7 7a"  # vsub.f32 s15, s15, s14
        " 17 ee 90 1a"  # vmov r1, s15
        " 00 f5 ad 30"  # add.w r0, r0, #0x15a00 (displaced stock instruction)
    )
    return body + encode_thumb_branch(CAVE_OFFSET + len(body), SETTER_CONTINUE)


def patch_image(source: bytes) -> bytes:
    if digest(source) != SOURCE_SHA256:
        raise ValueError(f"unknown source sha256={digest(source)}")
    if source[SETTER_OFFSET:SETTER_CONTINUE] != SETTER_ORIGINAL:
        raise ValueError("display-speed source setter differs")
    if source[CAVE_OFFSET:CAVE_OFFSET + CAVE_SIZE] != bytes(CAVE_SIZE):
        raise ValueError("speed source code cave differs")

    for site in DISPLAY_CALL_SITES:
        expected = encode_thumb_branch(site, CONVERT_SPEED, link=True)
        if source[site:site + 4] != expected:
            raise ValueError(f"display call site is not stock: 0x{site:x}")
    safety_call = encode_thumb_branch(SAFETY_CALL_SITE, CONVERT_SPEED, link=True)
    if source[SAFETY_CALL_SITE:SAFETY_CALL_SITE + 4] != safety_call:
        raise ValueError("overspeed call site differs")

    result = bytearray(source)
    wrapper = source_wrapper()
    if len(wrapper) > CAVE_SIZE:
        raise AssertionError("speed source wrapper exceeds the approved cave")
    result[CAVE_OFFSET:CAVE_OFFSET + len(wrapper)] = wrapper
    result[SETTER_OFFSET:SETTER_CONTINUE] = encode_thumb_branch(
        SETTER_OFFSET, CAVE_OFFSET
    )

    changed = {
        index for index, (before, after) in enumerate(zip(source, result))
        if before != after
    }
    allowed = set(range(SETTER_OFFSET, SETTER_CONTINUE))
    allowed.update(range(CAVE_OFFSET, CAVE_OFFSET + len(wrapper)))
    if not changed or not changed.issubset(allowed):
        raise AssertionError("change outside the central speed-source windows")
    for site in DISPLAY_CALL_SITES:
        expected = encode_thumb_branch(site, CONVERT_SPEED, link=True)
        if result[site:site + 4] != expected:
            raise AssertionError(f"display call site changed: 0x{site:x}")
    if result[SAFETY_CALL_SITE:SAFETY_CALL_SITE + 4] != safety_call:
        raise AssertionError("overspeed call site changed")
    if len(result) != len(source):
        raise AssertionError("native library size changed")
    return bytes(result)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("source", type=Path)
    parser.add_argument("destination", type=Path)
    args = parser.parse_args()
    if args.destination.exists():
        raise SystemExit("destination already exists")
    source = args.source.read_bytes()
    result = patch_image(source)
    temporary = args.destination.with_name(args.destination.name + ".new")
    temporary.write_bytes(result)
    temporary.replace(args.destination)
    print(f"source_sha256={digest(source)}")
    print(f"output_sha256={digest(result)}")
    print(f"setter_branch=0x{SETTER_OFFSET:x}:{result[SETTER_OFFSET:SETTER_CONTINUE].hex()}")
    print(f"wrapper=0x{CAVE_OFFSET:x}:{source_wrapper().hex()}")
    print("display_calls=STOCK:" + ",".join(f"0x{site:x}" for site in DISPLAY_CALL_SITES))
    print(f"overspeed_call=STOCK:0x{SAFETY_CALL_SITE:x}")


if __name__ == "__main__":
    main()