import time
from machine import Pin, I2C

class FT6336U:
    ADDR = 0x38
    
    # Registers
    REG_DEV_MODE = 0x00
    REG_TD_STATUS = 0x02
    REG_P1_XH = 0x03
    REG_P1_XL = 0x04
    REG_P1_YH = 0x05
    REG_P1_YL = 0x06
    REG_CTRL = 0x86
    REG_CHIPID = 0xA3
    REG_G_MODE = 0xA4
    
    # Modes
    CTRL_KEEP_ACTIVE = 0x00
    G_MODE_TRIGGER = 0x01
    
    def __init__(self, i2c, rst_pin=28, int_pin=25, width=480, height=320, swap_xy=True, invert_x=True, invert_y=False):
        self.i2c = i2c
        self.rst = Pin(rst_pin, Pin.OUT)
        self.int = Pin(int_pin, Pin.IN, Pin.PULL_UP)
        self.width = width
        self.height = height
        self.swap_xy = swap_xy
        self.invert_x = invert_x
        self.invert_y = invert_y
        self.initialized = False
        
        self.reset()
        self.init_chip()

    def reset(self):
        self.rst(0)
        time.sleep_ms(10)
        self.rst(1)
        time.sleep_ms(300) # Wait for chip to wake up

    def read_reg(self, reg, n=1):
        try:
            return self.i2c.readfrom_mem(self.ADDR, reg, n)
        except Exception as e:
            print(f"I2C read failed at reg 0x{reg:02X}: {e}")
            return None

    def write_reg(self, reg, val):
        try:
            self.i2c.writeto_mem(self.ADDR, reg, bytearray([val]))
            return True
        except Exception as e:
            print(f"I2C write failed at reg 0x{reg:02X}: {e}")
            return False

    def init_chip(self):
        # 1. Read Chip ID
        chip_id = self.read_reg(self.REG_CHIPID)
        if chip_id is None or chip_id[0] != 0x64:
            # Retry once
            time.sleep_ms(100)
            chip_id = self.read_reg(self.REG_CHIPID)
            if chip_id is None or chip_id[0] != 0x64:
                print(f"FT6336U error: Invalid Chip ID (got {chip_id[0] if chip_id else None}, expected 0x64)")
                self.initialized = False
                return False
                
        print(f"FT6336U touch controller detected (Chip ID: 0x{chip_id[0]:02X})")
        
        # 2. Configure operating mode (0 = Normal Mode)
        self.write_reg(self.REG_DEV_MODE, 0x00)
        
        # 3. Configure CTRL mode (0 = Keep Active)
        self.write_reg(self.REG_CTRL, self.CTRL_KEEP_ACTIVE)
        
        self.initialized = True
        return True

    def is_touched(self):
        """Returns True if screen is touched by checking the TD_STATUS register and clearing coordinates."""
        if not self.initialized:
            return False
        td_status = self.read_reg(self.REG_TD_STATUS)
        if td_status is None:
            return False
        touch_count = td_status[0] & 0x0F
        if touch_count > 0 and touch_count < 3:
            # Read P1 coordinates to clear register/interrupt state on the chip
            self.read_reg(self.REG_P1_XH, 6)
            return True
        return False

    def read_touch(self):
        """Reads touch point coordinates.
        Returns:
            (x, y) coordinates of first touch point, or None if no touch.
        """
        if not self.initialized:
            return None
            
        # Read TD_STATUS to check if there is an active touch
        td_status = self.read_reg(self.REG_TD_STATUS)
        if td_status is None:
            return None
            
        touch_count = td_status[0] & 0x0F
        if touch_count == 0:
            return None
            
        # Read P1 coordinates (6 bytes starting at P1_XH)
        buf = self.read_reg(self.REG_P1_XH, 6)
        if buf is None or len(buf) < 6:
            return None
            
        raw_x = ((buf[0] & 0x0F) << 8) | buf[1]
        raw_y = ((buf[2] & 0x0F) << 8) | buf[3]
        
        # Apply coordinate transformations to map touch coordinate space to screen space
        if self.swap_xy:
            raw_x, raw_y = raw_y, raw_x
            
        if self.invert_x:
            raw_x = self.width - 1 - raw_x
            
        if self.invert_y:
            raw_y = self.height - 1 - raw_y
            
        # Clamp coordinates to screen boundaries
        x = max(0, min(self.width - 1, raw_x))
        y = max(0, min(self.height - 1, raw_y))
        
        return (x, y)