#!/usr/bin/env python3
"""Read-only H618 PI-bank snapshot for no-load safety verification."""

from __future__ import annotations

import json
import mmap
import os
import struct
from pathlib import Path

PIO_BASE = 0x0300B000
PIO_LENGTH = 0x400
PI_OFFSET = 0x120
DATA_OFFSET = 0x10
OWNED_BITS = tuple([*range(0, 7), *range(9, 16)])
RESERVED_BITS = (7, 8, 16)


def read_u32(mapped: mmap.mmap, offset: int) -> int:
    return struct.unpack_from("<I", mapped, offset)[0]


def mux_mode(mapped: mmap.mmap, pin: int) -> int:
    register = PI_OFFSET + (pin // 8) * 4
    return (read_u32(mapped, register) >> ((pin % 8) * 4)) & 0xF


def device_tree_contract() -> dict:
    node = Path("/proc/device-tree/soc/pinctrl@300b000")
    compatible = (node / "compatible").read_bytes().rstrip(b"\0").decode("ascii")
    values = struct.unpack(">II", (node / "reg").read_bytes())
    if compatible != "allwinner,sun50i-h616-pinctrl" or values != (PIO_BASE, PIO_LENGTH):
        raise RuntimeError(f"unexpected pinctrl contract: {compatible=} {values=}")
    return {"compatible": compatible, "base": f"0x{values[0]:08X}", "length": values[1]}


def main() -> None:
    contract = device_tree_contract()
    descriptor = os.open("/dev/mem", os.O_RDONLY | os.O_SYNC)
    try:
        with mmap.mmap(descriptor, PIO_LENGTH, access=mmap.ACCESS_READ, offset=PIO_BASE) as mapped:
            data = read_u32(mapped, PI_OFFSET + DATA_OFFSET)
            modes = {f"PI{pin}": mux_mode(mapped, pin) for pin in range(17)}
    finally:
        os.close(descriptor)
    owned_mask = sum(1 << bit for bit in OWNED_BITS)
    reserved_mask = sum(1 << bit for bit in RESERVED_BITS)
    print(json.dumps({
        "device_tree": contract,
        "data_register": f"0x{data:08X}",
        "owned_data_bits": f"0x{data & owned_mask:08X}",
        "reserved_data_bits": f"0x{data & reserved_mask:08X}",
        "modes": modes,
    }, sort_keys=True))


if __name__ == "__main__":
    main()