from machine import Pin, SPI
import time
import ili9341
import glcdfont
import xpt2046

# ---------------- SPI ----------------
spi = SPI(
    2,
    baudrate=40000000,
    sck=Pin(18),
    mosi=Pin(23),
    miso=Pin(19)
)

# ---------------- DISPLAY ----------------
display = ili9341.ILI9341(
    spi,
    cs=Pin(15),
    dc=Pin(2),
    rst=Pin(4),
    width=240,
    height=320
)

# ---------------- TOUCH ----------------
touch = xpt2046.XPT2046(
    spi,
    cs=Pin(13),
    irq=Pin(12)
)

# ---------------- GLOBAL STATE ----------------
current_student = None
task_done = False

# ---------------- UI HELPERS ----------------
def draw_button(x, y, w, h, text, color):
    display.fill_rectangle(x, y, w, h, color)
    display.text(glcdfont, text, x + 10, y + 12, 0xFFFF)

def is_touched(x, y, w, h, tx, ty):
    return x < tx < (x + w) and y < ty < (y + h)

# ---------------- LOGIN SCREEN ----------------
def login_screen():
    global current_student
    display.fill(0)

    display.text(glcdfont, "SMART STUDY DESK", 20, 20, 0xFFFF)
    display.text(glcdfont, "Select Student", 40, 60, 0xFFFF)

    draw_button(40, 120, 160, 40, "Student A", 0x07E0)
    draw_button(40, 180, 160, 40, "Student B", 0x001F)

    while current_student is None:
        if touch.touched():
            tx, ty = touch.get_touch()
            if is_touched(40, 120, 160, 40, tx, ty):
                current_student = "Student A"
            elif is_touched(40, 180, 160, 40, tx, ty):
                current_student = "Student B"
            time.sleep(0.3)

# ---------------- TASK DASHBOARD ----------------
def task_screen():
    global task_done
    display.fill(0)

    display.text(glcdfont, "Welcome:", 20, 20, 0xFFFF)
    display.text(glcdfont, current_student, 20, 40, 0xFFFF)

    display.text(glcdfont, "Task: Study 30 min", 20, 80, 0xFFFF)

    draw_button(30, 140, 80, 40, "✓ DONE", 0x07E0)
    draw_button(130, 140, 80, 40, "✗ FAIL", 0xF800)

    while True:
        if touch.touched():
            tx, ty = touch.get_touch()
            if is_touched(30, 140, 80, 40, tx, ty):
                task_done = True
                show_result(True)
                break
            elif is_touched(130, 140, 80, 40, tx, ty):
                task_done = False
                show_result(False)
                break
            time.sleep(0.3)

# ... (Keep your existing imports and setup) ...

def show_result(success):
    display.fill(0)
    msg = "Task Completed" if success else "Task Incomplete"
    color = 0x07E0 if success else 0xF800
    display.text(glcdfont, msg, 40, 140, color)
    time.sleep(2)

# ---------------- IMPROVED MAIN FLOW ----------------
def main():
    login_screen()
    
    # Use a loop instead of calling functions inside each other
    while True:
        task_screen() 
        # Inside task_screen, when a button is pressed, 
        # it returns the success state instead of calling show_result