from __future__ import annotations

import threading
import time
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_I3C_BUS_INIT_NACK_RSTDAA,
    FW_INTERFACE_ALREADY_INITIALIZED,
    I2cTransferRate,
    I3cDriveStrength,
    I3cOpenDrainTransferRate,
    I3cPushPullRate,
    I3cTransferFlags,
    NotificationType,
    SUCCESS,
)
from example_helpers import print_device_info, print_hex_payload, assert_status


def main() -> None:
    sync = {
        "hot_join_received": False,
        "dynamic_address": 0,
        "ibi_count": 0,
        "last_ibi_payload": [],
    }
    sync_condition = threading.Condition()

    with CosmicPySDK(timeout_ms=2000) as device:
        def on_notification(notification) -> None:
            if notification.type == NotificationType.NOTIF_I3C_HJ and notification.i3c_hot_join is not None:
                with sync_condition:
                    sync["hot_join_received"] = True
                    sync["dynamic_address"] = notification.i3c_hot_join.dynamic_address
                    sync_condition.notify_all()
                print(f"Hot Join Notification: dynamic address 0x{notification.i3c_hot_join.dynamic_address:02X}")
                return

            if notification.type == NotificationType.NOTIF_I3C_IBI and notification.i3c_ibi is not None:
                with sync_condition:
                    sync["ibi_count"] += 1
                    sync["last_ibi_payload"] = list(notification.i3c_ibi.payload)
                    sync_condition.notify_all()
                print(
                    f"IBI Notification: address 0x{notification.i3c_ibi.address:02X} "
                    f"payloadLength {len(notification.i3c_ibi.payload)}"
                )

        device.on_notification(on_notification)

        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.i3c_set_bus_voltage(1200)
        assert_status(status, "i3c_set_bus_voltage")

        status, _ = device.i3c_controller_init(
            I3cPushPullRate.PUSH_PULL_10_MHz_DC_40,
            I3cOpenDrainTransferRate.OPEN_DRAIN_400_KHz,
            I2cTransferRate.FAST_MODE,
            I3cDriveStrength.FAST,
        )
        assert_status(status, "i3c_controller_init", (SUCCESS, FW_INTERFACE_ALREADY_INITIALIZED))

        status, _ = device.i3c_init_bus()
        assert_status(status, "i3c_init_bus", (SUCCESS, FW_I3C_BUS_INIT_NACK_RSTDAA))

        print("Power the PIC18F16Q20 target board now, waiting for Hot-Join...")
        with sync_condition:
            if not sync_condition.wait_for(lambda: sync["hot_join_received"], timeout=30.0):
                raise RuntimeError("Timed out waiting for Hot-Join notification.")

        dynamic_address = sync["dynamic_address"]
        print(f"PIC dynamic address: 0x{dynamic_address:02X}")
        time.sleep(0.25)

        status, _ = device.i3c_sdr_write(dynamic_address, [], [0x40, 0xA0, 0x00, 0x00, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5], I3cTransferFlags.I3C_XFER_NONE)
        assert_status(status, "i3c_sdr_write(I2C write)")
        time.sleep(0.25)

        status, _ = device.i3c_sdr_write(dynamic_address, [], [0x40, 0xA0, 0x00, 0x00], I3cTransferFlags.I3C_XFER_NONE)
        assert_status(status, "i3c_sdr_write(I2C pointer reset)")
        time.sleep(0.25)

        status, _ = device.i3c_sdr_write(dynamic_address, [], [0x20, 0xA0, 0x05], I3cTransferFlags.I3C_XFER_NONE)
        assert_status(status, "i3c_sdr_write(I2C read request)")
        expected_ibi_count = sync["ibi_count"] + 1
        with sync_condition:
            if not sync_condition.wait_for(lambda: sync["ibi_count"] >= expected_ibi_count, timeout=5.0):
                raise RuntimeError("Timed out waiting for I2C-read IBI response.")
            print_hex_payload("Last IBI payload:", sync["last_ibi_payload"])

        time.sleep(0.25)

        status, _ = device.i3c_sdr_write(dynamic_address, [], [0x60, 0x03, 0x9F], I3cTransferFlags.I3C_XFER_NONE)
        assert_status(status, "i3c_sdr_write(SPI JEDEC read request)")
        expected_ibi_count = sync["ibi_count"] + 1
        with sync_condition:
            if not sync_condition.wait_for(lambda: sync["ibi_count"] >= expected_ibi_count, timeout=5.0):
                raise RuntimeError("Timed out waiting for SPI JEDEC IBI response.")
            print_hex_payload("Last IBI payload:", sync["last_ibi_payload"])


if __name__ == "__main__":
    main()
