import os
import sys
import time
import spidev
import numpy as np
import lgpio
from smbus2 import SMBus
from PIL import Image, ImageDraw, ImageFont

LCD_WIDTH  = 320
LCD_HEIGHT = 240
MENU_W     = 80

# Pin definitions for Raspberry Pi GPIO(using BCM numbering)
LCD_RST_PIN = 22
LCD_DC_PIN  = 25
LCD_BL_PIN  = 24
TP_RST_PIN  = 27

h = lgpio.gpiochip_open(0)

class LCD_2Inch:
    def __init__(self):
        lgpio.gpio_claim_output(h, LCD_DC_PIN, 0)
        lgpio.gpio_claim_output(h, LCD_RST_PIN, 0)
        lgpio.gpio_claim_output(h, LCD_BL_PIN, 0)
        
        lgpio.tx_pwm(h, LCD_BL_PIN, 1000, 100) 

        self.spi = spidev.SpiDev()
        self.spi.open(0, 0)
        self.spi.max_speed_hz = 8000000 
        self.spi.mode = 0b00
        self.lcd_init()

    def write_cmd(self, cmd):
        lgpio.gpio_write(h, LCD_DC_PIN, 0)
        self.spi.writebytes([cmd])

    def write_data(self, value):
        lgpio.gpio_write(h, LCD_DC_PIN, 1)
        if isinstance(value, int):
            self.spi.writebytes([value])
        else:
            self.spi.writebytes(value)

    def reset(self):
        lgpio.gpio_write(h, LCD_RST_PIN, 1); time.sleep(0.01)
        lgpio.gpio_write(h, LCD_RST_PIN, 0); time.sleep(0.01)
        lgpio.gpio_write(h, LCD_RST_PIN, 1); time.sleep(0.1)

    def lcd_init(self):
        self.reset()
        self.write_cmd(0x36) 
        self.write_data(0x00)
        self.write_cmd(0x3A)
        self.write_data(0x05) 
        self.write_cmd(0x21) 
        self.write_cmd(0xB2)
        self.write_data([0x0C, 0x0C, 0x00, 0x33, 0x33])
        self.write_cmd(0xB7) 
        self.write_data(0x35)
        self.write_cmd(0xBB) 
        self.write_data(0x1F) 
        self.write_cmd(0xC0) 
        self.write_data(0x2C)
        self.write_cmd(0xC2) 
        self.write_data(0x01)
        self.write_cmd(0xC3) 
        self.write_data(0x12)
        self.write_cmd(0xC4) 
        self.write_data(0x20)
        self.write_cmd(0xC6) 
        self.write_data(0x0F)
        self.write_cmd(0xD0) 
        self.write_data([0xA4, 0xA1])
        self.write_cmd(0xE0) 
        self.write_data([0xD0,0x08,0x11,0x08,0x0C,0x15,0x39,0x33,0x50,0x36,0x13,0x14,0x29,0x2D])
        self.write_cmd(0xE1) 
        self.write_data([0xD0,0x08,0x10,0x08,0x06,0x06,0x39,0x44,0x51,0x0B,0x16,0x14,0x2F,0x31])
        self.write_cmd(0x21)
        self.write_cmd(0x11)
        time.sleep(0.12)
        self.write_cmd(0x29)

    def display(self, img):
        if img.size != (LCD_WIDTH, LCD_HEIGHT): 
            img = img.resize((LCD_WIDTH, LCD_HEIGHT))
        self.write_cmd(0x36)
        self.write_data(0x70)
        self.write_cmd(0x2A)
        self.write_data([0, 0, (LCD_WIDTH-1)>>8, (LCD_WIDTH-1)&0xFF])
        self.write_cmd(0x2B)
        self.write_data([0, 0, (LCD_HEIGHT-1)>>8, (LCD_HEIGHT-1)&0xFF])
        self.write_cmd(0x2C)
        
        lgpio.gpio_write(h, LCD_DC_PIN, 1)
        arr = np.asarray(img)
        
        BLOCK_H = 20
        for y in range(0, LCD_HEIGHT, BLOCK_H):
            chunk = arr[y:y+BLOCK_H]
            ch_h, ch_w, _ = chunk.shape
            pixel = np.zeros((ch_h, ch_w, 2), dtype=np.uint8)
            pixel[..., 0] = (chunk[..., 0] & 0xF8) | (chunk[..., 1] >> 5)
            pixel[..., 1] = ((chunk[..., 1] << 3) & 0xE0) | (chunk[..., 2] >> 3)
            b_data = pixel.tobytes()
            for i in range(0, len(b_data), 4096):
                self.spi.writebytes(list(b_data[i:i+4096]))

    def draw_point_fast(self, x, y, color_rgb):
        if x >= LCD_WIDTH-2 or y >= LCD_HEIGHT-2: return
        r, g, b = color_rgb
        hi = (r & 0xF8) | (g >> 5); lo = ((g << 3) & 0xE0) | (b >> 3)
        self.write_cmd(0x2A); self.write_data([x>>8, x&0xFF, (x+2)>>8, (x+2)&0xFF])
        self.write_cmd(0x2B); self.write_data([y>>8, y&0xFF, (y+2)>>8, (y+2)&0xFF])
        self.write_cmd(0x2C)
        lgpio.gpio_write(h, LCD_DC_PIN, 1)
        self.spi.writebytes([hi, lo] * 9)

class Touch_CST816:
    def __init__(self):
        lgpio.gpio_claim_output(h, TP_RST_PIN, 1)
        print("Touch: Resetting...")
        lgpio.gpio_write(h, TP_RST_PIN, 1)  # High
        time.sleep(0.05)
        lgpio.gpio_write(h, TP_RST_PIN, 0)  # Low (Reset)
        time.sleep(0.02) 
        lgpio.gpio_write(h, TP_RST_PIN, 1)  # High 
        time.sleep(0.3) 
        
        try:
            self.bus = SMBus(1)  
            self.bus.read_byte(0x15)
            print("Touch:I2C connection successful (0x15)")
        except Exception as e:
            print(f"\n!!! Touch Error: I2C connection failed !!!")
            sys.exit(1)
            
    def get_coord(self):
        try:
            data = self.bus.read_i2c_block_data(0x15, 0x01, 7)
            fingers = data[2] 
            if fingers > 0: 
                raw_y = ((data[3] & 0x0F) << 8) | data[4]
                raw_x = ((data[5] & 0x0F) << 8) | data[6]
               
                x = raw_x
                y = raw_y
                return x, y
        except:
            pass
        return None

def main():
    print("Initializing...")
    lcd = LCD_2Inch()
    touch = Touch_CST816()
    
    demo_img = Image.new("RGB", (LCD_WIDTH, LCD_HEIGHT), "WHITE")
    draw_demo = ImageDraw.Draw(demo_img)
    
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 16)
    except:
        font = ImageFont.load_default()
    try:
        bg = Image.open('fruits.jpg').resize((LCD_WIDTH, LCD_HEIGHT))
        demo_img.paste(bg, (0,0))
    except:
        draw_demo.rectangle([0, 0, LCD_WIDTH, LCD_HEIGHT], fill="LIGHTBLUE")

    lcd.display(demo_img)
    time.sleep(2)
    
    draw_demo.rectangle([0, 0, LCD_WIDTH, LCD_HEIGHT], fill="WHITE")
    draw_demo.line([20, 80, 80, 80], fill="red", width=2)

    draw_demo.rectangle([20, 10, 70, 60], fill="WHITE", outline="BLUE") 
    draw_demo.rectangle([85, 10, 130, 60], outline="BLUE", width=3) 
    draw_demo.arc((160, 30, 210, 80),0,360,fill=(255,0,255)) 
    draw_demo.ellipse([240, 30, 290, 80], fill="YELLOW", outline="ORANGE") 

    draw_demo.text((20, 120), "SEENGREAT", fill="BLUE", font=font)
    draw_demo.text((20, 160), "2 inch LCD", fill="PINK", font=font)
    draw_demo.text((20, 200), "ab CD 3210", fill="RED", font=font)

    lcd.display(demo_img)
    time.sleep(2)
    
    canvas = Image.new("RGB", (LCD_WIDTH, LCD_HEIGHT), "WHITE")
    draw = ImageDraw.Draw(canvas)
    draw.rectangle([0, 0, MENU_W, LCD_HEIGHT], fill="WHITE")
    buttons = [
        (5, 'CLR', (0,0,255)), (44, 'RED', (255,0,0)), (83, 'GRN', (0,255,0)),
        (122, 'BLU', (0,0,255)), (161, 'BLK', (0,0,0)), (200, 'YEL', (255,255,0))
    ]
    for y, t, c in buttons:
        draw.rectangle([5, y, 75, y+34], fill=c, outline="BLACK")
        if t == 'CLR': draw.text((20, y+8), "CLR", fill="WHITE", font=font)
    draw.line([MENU_W, 0, MENU_W, LCD_HEIGHT], fill="BLACK", width=2)

    lcd.display(canvas)
    
    curr_color = (0, 0, 0)
    last_act = 0
    
    print("Touch Loop Started...")
    try:
        while True:

             pos = touch.get_coord()

             if pos:
                x, y = pos
                
                if x < MENU_W:
                    if time.time() - last_act > 0.3:
                        
                        if 5 <= y <= 39:
                            draw.rectangle([MENU_W+2, 0, LCD_WIDTH, LCD_HEIGHT], fill="WHITE")
                            lcd.display(canvas)
                        elif 44 <= y <= 78: curr_color = (255, 0, 0)
                        elif 83 <= y <= 117: curr_color = (0, 255, 0)
                        elif 122 <= y <= 156: curr_color = (0, 0, 255)
                        elif 161 <= y <= 195: curr_color = (0, 0, 0)
                        elif 200 <= y <= 234: curr_color = (255, 255, 0)
                        last_act = time.time()
                else:
                    lcd.draw_point_fast(x, y, curr_color)
                    draw.point([x,y], fill=curr_color)
             time.sleep(0.001)
                    
    except KeyboardInterrupt:
        print("\nExit")
    finally:

        lgpio.gpiochip_close(h)

if __name__ == "__main__":
    main()

