import queue
import time
from typing import List, Optional
from iowp.special import SpecialModeFunction
from iowp.devices import IowDevice

class I2C(SpecialModeFunction):
    """
    Implementation of the I2C special mode function.
    Matches de.wagner_ibw.iow.i2c.I2C.
    """
    I2C_SPEED_100KHZ = 0
    I2C_SPEED_400KHZ = 1
    I2C_SPEED_50KHZ  = 2

    GENERATE_START = 0x80
    GENERATE_STOP  = 0x40

    def __init__(self, speed: int = I2C_SPEED_100KHZ):
        super().__init__()
        self.i2c_speed = speed
        self._write_ack_queue = queue.Queue()
        self._read_queue = queue.Queue()

    def get_name(self) -> str:
        return "I2C"

    def get_special_mode_function_id(self) -> int:
        return self.SMF_I2C_ID

    def get_report_ids(self) -> List[int]:
        return [0x02, 0x03]

    def report_received(self, read_buffer: List[int]):
        if read_buffer[0] == 0x02:
            self._write_ack_queue.put(read_buffer)
        elif read_buffer[0] == 0x03:
            self._read_queue.put(read_buffer)

    def get_enable_report(self) -> List[int]:
        return [0x01, 0x01, self.i2c_speed]

    def get_disable_report(self) -> List[int]:
        return [0x01, 0x00]

    def get_iow_special_bits(self, product_id: int) -> List[int]:
        if product_id == IowDevice.IOW24ID:
            return [0x06]
        elif product_id == IowDevice.IOW56ID:
            return [0, 0xA0]
        else: # IOW40 / IOW28 / IOW100
            return [0xC0]



    def write_i2c(self, addr: int, data: List[int]):
        """
        Writes data to an I2C device with address addr.
        Note: addr should include the R/W bit if following Java style, 
              but usually it's just the 7-bit address.
              Java code seems to take the data as is.
        """
        if not self.iow:
            raise RuntimeError("I2C not attached to a device")

        # Simplified chunking logic (similar to Java writeI2C)
        length = len(data)
        pointer = 0
        first = True
        
        while pointer < length:
            remain = length - pointer
            part_len = min(remain, 6) # Max 6 bytes per report for IOW24/40/28
            if self.iow.product_id == IowDevice.IOW56ID:
                part_len = min(remain, 62)

            flags = part_len
            if first:
                flags |= self.GENERATE_START
                first = False
            if (pointer + part_len) == length:
                flags |= self.GENERATE_STOP

            report = [0x02, flags] + data[pointer:pointer+part_len]
            self.iow.write_report(1, report)
            pointer += part_len

            # Wait for ACK
            try:
                ack_buf = self._write_ack_queue.get(timeout=2.0)
                # Check for error bit (0x80 in flags)
                if ack_buf[1] & 0x80:
                    raise IOError(f"I2C Write Error: 0x{ack_buf[1]:02X}")
            except queue.Empty:
                raise IOError("I2C Write Timeout (ACK not received)")

    def read_i2c(self, addr: int, count: int) -> List[int]:
        """Reads count bytes from an I2C device."""
        if not self.iow:
            raise RuntimeError("I2C not attached to a device")

        # Start combined write/read report
        # Report ID 3, Count, Data...
        # For a simple read, we just send Report ID 3 with the count.
        report = [0x03, count]
        self.iow.write_report(1, report)

        result = []
        remain = count
        while remain > 0:
            try:
                read_buf = self._read_queue.get(timeout=2.0)
                if read_buf[1] & 0x80:
                    raise IOError(f"I2C Read Error: 0x{read_buf[1]:02X}")
                
                got = read_buf[1] & 0x3F # Mask out flags
                if got == 0:
                    raise IOError("I2C Read returned 0 bytes")
                
                result.extend(read_buf[2:2+got])
                remain -= got
            except queue.Empty:
                raise IOError("I2C Read Timeout")
        
        return result
