import struct
from driver.protocol import I2CPCClient, time
from driver.patches import *
import time

from typing import TypeVar, ParamSpec

P = ParamSpec("P")
T = TypeVar("T")

class BaseTEF668X:
    def __init__(self, p: I2CPCClient, address: int = 0x64) -> None:
        self.p = p
        p.set_baudrate(921600)
        self.p.set_clock(100_000)
        self.address = address
    def close(self):
        self.p.set_baudrate(115200)
        self.p.close()
    def _reset(self):
        self.p.write_i2c(self.address, b"\x1e\x5a\x01\x5a\x5a")
    def _write_firmware(self, patch: bytes | list[int], patch_lut: bytes | list[int]):
        self.p.set_clock(400_000)
        self.p.write_i2c(self.address, b"\x1c\x00\x00")
        self.p.write_i2c(self.address, b"\x1c\x00\x74")
        def send_patch(_patch: bytes):
            for i in range(0, len(_patch), 20):
                data = _patch[i:i+20]
                if self.p.write_i2c(self.address, b"\x1b" + data)[0] != 0: raise Exception
        send_patch(bytes(patch))
        self.p.write_i2c(self.address, b"\x1c\x00\x00")
        self.p.write_i2c(self.address, b"\x1c\x00\x75")
        send_patch(bytes(patch_lut))
        self.p.write_i2c(self.address, b"\x1c\x00\x00")
        self.p.set_clock(100_000)
    
    def APPL_Get_Operation_Status(self):
        data = self.p.write_read_i2c(self.address, b"\x40\x80\x01", 2)
        start = time.monotonic()
        while data[0] != 0:
            data = self.p.write_read_i2c(self.address, b"\x40\x80\x01", 2)
            time.sleep(0.01)
            if (time.monotonic() - start) > 1: raise TimeoutError("Could not get status")
        return data[-1]

    def APPL_Set_ReferenceClock(self, clock: int, type_clock: bool):
        return self.p.write_i2c(self.address, b"\x40\x04\x01" + struct.pack(">IH", clock, type_clock))
    def APPL_Activate(self): return self.p.write_i2c(self.address, b"\x40\x05\x01\x00\x01")

    def init(self, patch = tef_102_patch, patch_lut = tef_102_patch_lut, clock: int = 9216000):
        self._reset()
        while self.APPL_Get_Operation_Status() != 0: time.sleep(0.025) # TODO: timeouts
        self._write_firmware(patch, patch_lut)
        self.p.write_i2c(self.address, b"\x14\x00\x01")
        while self.APPL_Get_Operation_Status() != 1: time.sleep(0.025)
        if clock != 9216000: self.APPL_Set_ReferenceClock(clock, False)
        self.APPL_Activate()
        while self.APPL_Get_Operation_Status() != 2: time.sleep(0.025)

    def __enter__(self): return self

    def __exit__(self, exc_type, exc_val, exc_tb): self.close()