"""
ESP32-S3 TFT Display Test Script
MicroPython code for testing TFT displays
Supports ST7789, ILI9341, and other common drivers

Author: Arvind Patil
AI Centre Nandurbar
"""

from machine import Pin, SPI
import time

# ==== Display Configuration ====
# Common ESP32-S3 SPI pins for TFT
SPI_SCK = 12   # Clock
SPI_MOSI = 11  # Data (MOSI/SDA)
SPI_MISO = 13  # Not used for display but needed for SPI init
DC_PIN = 8     # Data/Command
CS_PIN = 10    # Chip Select
RST_PIN = 9    # Reset
BL_PIN = 14    # Backlight (optional)

# Display dimensions (adjust as needed)
WIDTH = 240
HEIGHT = 320

# ==== Color Definitions (RGB565) ====
BLACK = 0x0000
WHITE = 0xFFFF
RED = 0xF800
GREEN = 0x07E0
BLUE = 0x001F
YELLOW = 0xFFE0
CYAN = 0x07FF
MAGENTA = 0xF81F

class TFT_Test:
    def __init__(self):
        """Initialize TFT display"""
        print("🔧 Initializing ESP32-S3 TFT Display...")
        
        # Initialize SPI
        self.spi = SPI(
            1,
            baudrate=40000000,  # 40MHz
            polarity=0,
            phase=0,
            sck=Pin(SPI_SCK),
            mosi=Pin(SPI_MOSI),
            miso=Pin(SPI_MISO)
        )
        
        # Initialize control pins
        self.dc = Pin(DC_PIN, Pin.OUT)
        self.cs = Pin(CS_PIN, Pin.OUT)
        self.rst = Pin(RST_PIN, Pin.OUT)
        
        # Backlight (if available)
        try:
            self.bl = Pin(BL_PIN, Pin.OUT)
            self.bl.value(1)  # Turn on backlight
            print("✅ Backlight enabled on pin", BL_PIN)
        except:
            self.bl = None
            print("⚠️  No backlight pin configured")
        
        # Reset display
        self.reset()
        
        # Try to initialize display
        self.init_display()
        
        print("✅ Display initialized successfully!")
    
    def reset(self):
        """Hardware reset of display"""
        print("🔄 Resetting display...")
        self.rst.value(1)
        time.sleep_ms(5)
        self.rst.value(0)
        time.sleep_ms(20)
        self.rst.value(1)
        time.sleep_ms(150)
    
    def write_cmd(self, cmd):
        """Send command to display"""
        self.dc.value(0)  # Command mode
        self.cs.value(0)
        self.spi.write(bytearray([cmd]))
        self.cs.value(1)
    
    def write_data(self, data):
        """Send data to display"""
        self.dc.value(1)  # Data mode
        self.cs.value(0)
        if isinstance(data, int):
            self.spi.write(bytearray([data]))
        else:
            self.spi.write(bytearray(data))
        self.cs.value(1)
    
    def init_display(self):
        """Initialize display with common commands (ST7789/ILI9341)"""
        print("📱 Sending initialization commands...")
        
        # Software reset
        self.write_cmd(0x01)
        time.sleep_ms(150)
        
        # Sleep out
        self.write_cmd(0x11)
        time.sleep_ms(255)
        
        # Color mode - 16bit (RGB565)
        self.write_cmd(0x3A)
        self.write_data(0x55)
        
        # Memory access control (adjust rotation)
        self.write_cmd(0x36)
        self.write_data(0x00)  # Normal orientation
        
        # Display on
        self.write_cmd(0x29)
        time.sleep_ms(100)
        
        # Normal display mode
        self.write_cmd(0x13)
        time.sleep_ms(10)
    
    def set_window(self, x0, y0, x1, y1):
        """Set drawing window"""
        # Column address set
        self.write_cmd(0x2A)
        self.write_data([x0 >> 8, x0 & 0xFF, x1 >> 8, x1 & 0xFF])
        
        # Row address set
        self.write_cmd(0x2B)
        self.write_data([y0 >> 8, y0 & 0xFF, y1 >> 8, y1 & 0xFF])
        
        # Memory write
        self.write_cmd(0x2C)
    
    def fill_color(self, color):
        """Fill entire screen with color"""
        print(f"🎨 Filling screen with color 0x{color:04X}")
        
        self.set_window(0, 0, WIDTH-1, HEIGHT-1)
        
        # Create buffer with color
        buf = bytearray(WIDTH * 2)
        for i in range(0, len(buf), 2):
            buf[i] = color >> 8
            buf[i+1] = color & 0xFF
        
        # Write buffer for each row
        self.dc.value(1)
        self.cs.value(0)
        for _ in range(HEIGHT):
            self.spi.write(buf)
        self.cs.value(1)
        
        print("✅ Fill complete")
    
    def draw_pixel(self, x, y, color):
        """Draw single pixel"""
        if x < 0 or x >= WIDTH or y < 0 or y >= HEIGHT:
            return
        
        self.set_window(x, y, x, y)
        self.write_data([color >> 8, color & 0xFF])
    
    def draw_rect(self, x, y, w, h, color):
        """Draw filled rectangle"""
        print(f"📐 Drawing rectangle at ({x},{y}) size {w}x{h}")
        
        if x + w > WIDTH:
            w = WIDTH - x
        if y + h > HEIGHT:
            h = HEIGHT - y
        
        self.set_window(x, y, x+w-1, y+h-1)
        
        # Create buffer with color
        buf = bytearray(w * 2)
        for i in range(0, len(buf), 2):
            buf[i] = color >> 8
            buf[i+1] = color & 0xFF
        
        # Write buffer for each row
        self.dc.value(1)
        self.cs.value(0)
        for _ in range(h):
            self.spi.write(buf)
        self.cs.value(1)
    
    def test_colors(self):
        """Test display with different colors"""
        print("\n🌈 === COLOR TEST SEQUENCE ===")
        
        colors = [
            (RED, "RED"),
            (GREEN, "GREEN"),
            (BLUE, "BLUE"),
            (YELLOW, "YELLOW"),
            (CYAN, "CYAN"),
            (MAGENTA, "MAGENTA"),
            (WHITE, "WHITE"),
            (BLACK, "BLACK")
        ]
        
        for color, name in colors:
            print(f"Testing {name}...")
            self.fill_color(color)
            time.sleep(1)
    
    def test_gradients(self):
        """Test color gradients"""
        print("\n🌈 === GRADIENT TEST ===")
        
        # Horizontal RGB gradient
        print("Drawing horizontal RGB gradient...")
        for x in range(WIDTH):
            r = int((x / WIDTH) * 31) << 11
            g = int((x / WIDTH) * 63) << 5
            b = int((x / WIDTH) * 31)
            color = r | g | b
            
            for y in range(HEIGHT):
                self.draw_pixel(x, y, color)
        
        time.sleep(2)
    
    def test_patterns(self):
        """Test geometric patterns"""
        print("\n📐 === PATTERN TEST ===")
        
        # Clear screen
        self.fill_color(BLACK)
        time.sleep(0.5)
        
        # Draw colored rectangles
        print("Drawing colored rectangles...")
        rect_h = HEIGHT // 4
        self.draw_rect(0, 0, WIDTH, rect_h, RED)
        self.draw_rect(0, rect_h, WIDTH, rect_h, GREEN)
        self.draw_rect(0, rect_h*2, WIDTH, rect_h, BLUE)
        self.draw_rect(0, rect_h*3, WIDTH, rect_h, YELLOW)
        
        time.sleep(2)
        
        # Checkerboard pattern
        print("Drawing checkerboard pattern...")
        self.fill_color(BLACK)
        square_size = 20
        
        for y in range(0, HEIGHT, square_size):
            for x in range(0, WIDTH, square_size):
                if (x // square_size + y // square_size) % 2:
                    self.draw_rect(x, y, square_size, square_size, WHITE)
        
        time.sleep(2)
    
    def run_full_test(self):
        """Run complete test sequence"""
        print("\n" + "="*50)
        print("🚀 ESP32-S3 TFT DISPLAY FULL TEST")
        print("="*50)
        
        # Color test
        self.test_colors()
        
        # Pattern test
        self.test_patterns()
        
        # Gradient test (commented out as it's slow)
        #self.test_gradients()
        
        # Final screen
        self.fill_color(GREEN)
        
        print("\n" + "="*50)
        print("✅ TEST COMPLETE!")
        print("="*50)


# ==== Main execution ====
def main():
    """Main test function"""
    try:
        # Create and initialize display
        tft = TFT_Test()
        
        # Run full test
        tft.run_full_test()
        
        print("\n💡 TIP: Adjust pin numbers at the top if needed")
        print("📝 Display dimensions: {}x{}".format(WIDTH, HEIGHT))
        
    except Exception as e:
        print(f"❌ Error: {e}")
        import sys
        sys.print_exception(e)


if __name__ == "__main__":
    main()


# ==== Usage Instructions ====
"""
UPLOAD करण्यासाठी:
1. Thonny IDE उघडा
2. ESP32-S3 USB ने कनेक्ट करा
3. हा फाइल उघडा आणि "Run" दाबा
4. किंवा REPL मध्ये: import esp32_s3_tft_test

PIN CONFIGURATION:
- SCK/CLK: Pin 12
- MOSI/SDA: Pin 11
- DC: Pin 8
- CS: Pin 10
- RST: Pin 9
- BL (Backlight): Pin 14

तुमच्या TFT display च्या pins वेगळे असतील तर वरील numbers बदला.

TROUBLESHOOTING:
- Display काहीही दाखवत नसेल तर backlight pin चेक करा
- Colors चुकीचे असतील तर memory access control (0x36 command) adjust करा
- काम करत नसेल तर baudrate कमी करा (20000000)
"""