from __future__ import annotations

import random
import sys
from pathlib import Path


ROOT = next(parent for parent in Path(__file__).resolve().parents if (parent / "cosmic_driver").is_dir())
sys.path.insert(0, str(ROOT))

from cosmic_driver import (
    CosmicPySDK,
    FW_INTERFACE_ALREADY_INITIALIZED,
    SUCCESS,
    SpiBitOrder,
    SpiChipSelect,
    SpiChipSelectPolarity,
    SpiDataWidth,
    SpiMode,
)
from example_helpers import print_device_info, assert_status


def main() -> None:
    with CosmicPySDK(timeout_ms=2000) as device:
        if not device.open():
            raise RuntimeError("Failed to connect to a Binho device.")

        status, info = device.get_device_info()
        assert_status(status, "get_device_info")
        assert info is not None
        print_device_info(info)

        status, _ = device.spi_set_bus_voltage(3300)
        assert_status(status, "spi_set_bus_voltage")

        status, _ = device.spi_controller_init(
            SpiBitOrder.SPI_BIT_ORDER_MSB,
            SpiMode.SPI_MODE_0,
            SpiDataWidth.SPI_DATA_WIDTH_8_BITS,
            SpiChipSelect.SPI_CHIP_SELECT_0,
            SpiChipSelectPolarity.SPI_CS_ACTIVE_LOW,
            10_000_000,
        )
        assert_status(status, "spi_controller_init", (SUCCESS, FW_INTERFACE_ALREADY_INITIALIZED))

        status, jedec = device.spi_transfer([0x9F], 4)
        assert_status(status, "read JEDEC ID")
        assert jedec is not None
        if len(jedec) < 4 or jedec[1:4] != [0xEF, 0x40, 0x18]:
            raise RuntimeError(f"Unexpected JEDEC ID: {jedec}")

        status, _ = device.spi_transfer([0x06], 1)
        assert_status(status, "write enable")

        address = [0x00, 0x20, 0x00]
        status, _ = device.spi_transfer([0x20, *address], 4)
        assert_status(status, "sector erase")

        while True:
            status, register = device.spi_transfer([0x05], 2)
            assert_status(status, "read status register")
            assert register is not None
            if len(register) >= 2 and (register[1] & 0x01) == 0:
                break

        status, _ = device.spi_transfer([0x06], 1)
        assert_status(status, "write enable (again)")

        data_to_write = [random.randint(0, 255) for _ in range(random.randint(1, 256))]
        status, _ = device.spi_transfer([0x02, *address, *data_to_write], 4 + len(data_to_write))
        assert_status(status, "page program")

        while True:
            status, register = device.spi_transfer([0x05], 2)
            assert_status(status, "read status register (after program)")
            assert register is not None
            if len(register) >= 2 and (register[1] & 0x01) == 0:
                break

        status, readback = device.spi_transfer([0x03, *address], 4 + len(data_to_write))
        assert_status(status, "read data")
        assert readback is not None
        if readback[4:] != data_to_write:
            raise RuntimeError("Readback verification failed.")

        print(f"SPI flash write/read verification successful ({len(data_to_write)} bytes)")


if __name__ == "__main__":
    main()
