import machine
import utime
import math
import network
import urequests
from ili9341 import Display, color565

# --- 1. HARDWARE SETUP ---
spi = machine.SPI(0, baudrate=10000000, sck=machine.Pin(18), mosi=machine.Pin(19))
display = Display(spi, dc=machine.Pin(20), cs=machine.Pin(17), rst=machine.Pin(21), rotation=90)

row_pins, col_pins = [0, 1, 2, 3], [6, 7, 8, 9]
rows = [machine.Pin(p, machine.Pin.OUT) for p in row_pins]
cols = [machine.Pin(p, machine.Pin.IN, machine.Pin.PULL_DOWN) for p in col_pins]

# --- 2. SYSTEM STATE ---
key_map_normal = [['1','2','3','DEL'],['4','5','6','X'],['7','8','9','/'],['AC','0','sin','ANS']]
key_map_ai = [['a','b','c','DEL'],['d','e','f',' '],['g','h','i','?'],['AC','0','y','ANS']]

MODE_NORMAL, MODE_AI = 0, 1
current_mode = MODE_NORMAL
SECRET_CODE = "777"
input_buf = ""

# --- 3. UI ENGINE ---

def draw_ui():
    display.clear(0)
    if current_mode == MODE_AI:
        display.fill_rectangle(0, 0, 320, 22, color565(0, 120, 0))
        display.draw_text8x8(10, 7, "AI TERMINAL", 0)
        display.draw_hline(10, 165, 300, color565(50, 50, 50)) # Status Bar Line
    else:
        display.draw_text8x8(10, 7, "CASIO fx-991ES Plus", color565(100, 100, 100))
    
    display.draw_rectangle(10, 70, 300, 110, color565(255, 255, 0)) # Yellow Border
    icon_color = color565(150, 150, 150)
    display.draw_text8x8(225, 75, "D", icon_color) 
    display.draw_text8x8(250, 75, "Math", icon_color)

def update_status(msg):
    """Status bar at bottom - AI Mode Only"""
    if current_mode == MODE_AI:
        display.fill_rectangle(11, 166, 298, 13, 0)
        display.draw_text8x8(15, 170, msg.upper(), color565(180, 180, 180))

def update_input(text):
    """Normal Mode: Big/Tall input | AI Mode: Small/Stealth input"""
    display.fill_rectangle(15, 75, 205, 30, 0) # Clear input area
    if current_mode == MODE_AI:
        display.draw_text8x8(20, 80, "> " + text, color565(255, 255, 255))
    else:
        # Scale 2 (16x16) for authentic tall calculator input
        display.draw_text8x8(20, 78, text, color565(255, 255, 255), size=2)

def show_result(text):
    """Normal Mode: Giant Result | AI Mode: Multi-line Green with Status Bar"""
    if current_mode == MODE_AI:
        display.fill_rectangle(11, 105, 298, 60, 0) # Clear AI content zone
        if not text: return
        max_ch = 35
        lines = [text[i:i+max_ch] for i in range(0, len(text), max_ch)]
        for i, line in enumerate(reversed(lines[-4:])): 
            display.draw_text8x8(15, 152 - (i*10), line, color565(0, 255, 0))
    else:
        display.fill_rectangle(11, 115, 298, 63, 0) # Clear Math result zone
        if not text: return
        # Scale 3 (24x24) for the chunky, massive Casio answer look
        char_width = 24 
        x_pos = 295 - (len(text) * char_width)
        display.draw_text8x8(max(20, x_pos), 140, text, color565(255, 255, 255), size=3)

# --- 4. WIFI MANAGER ---

def scan_keys():
    active_map = key_map_normal if current_mode == MODE_NORMAL else key_map_ai
    for r_idx, rp in enumerate(rows):
        rp.value(1)
        for c_idx, cp in enumerate(cols):
            if cp.value() == 1:
                rp.value(0)
                return active_map[r_idx][c_idx]
        rp.value(0)
    return None

def wifi_setup_flow():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    update_status("SCANNING...")
    nets = wlan.scan()
    unique_nets = [n[0].decode('utf-8') for n in nets if n[0]]
    if not unique_nets: 
        update_status("NO NETS")
        return False

    idx = 0
    while True: # Network Selection Menu
        menu = "PICK:\n"
        for i in range(min(3, len(unique_nets))):
            ptr = ">" if i == idx else " "
            menu += f"{ptr}{unique_nets[i][:15]}\n"
        show_result(menu)
        
        k = scan_keys()
        if k == '1': idx = (idx - 1) % len(unique_nets)
        elif k == '7': idx = (idx + 1) % len(unique_nets)
        elif k == 'ANS': 
            target_ssid = unique_nets[idx]
            break
        elif k == 'AC': return False
        utime.sleep(0.15)

    # Password Input
    pass_buf = ""
    while True:
        show_result(f"PASS {target_ssid}:\n> " + pass_buf)
        k = scan_keys()
        if k:
            if k == 'AC': return False
            elif k == 'DEL': pass_buf = pass_buf[:-1]
            elif k == 'ANS': break
            else: pass_buf += k
            utime.sleep(0.2)
    
    update_status("CONNECTING...")
    wlan.connect(target_ssid, pass_buf)
    for _ in range(10):
        if wlan.isconnected(): return True
        utime.sleep(1)
    return False

# --- 5. MAIN LOOP ---
draw_ui()
last_k = ""

while True:
    k = scan_keys()
    if k and k != last_k:
        if k == 'AC':
            input_buf = ""
            current_mode = MODE_NORMAL
            draw_ui()
        elif k == 'DEL':
            input_buf = input_buf[:-1]
            update_input(input_buf)
        elif k == 'ANS':
            if input_buf == SECRET_CODE:
                current_mode = MODE_AI
                input_buf = ""
                draw_ui()
                if wifi_setup_flow(): update_status("ONLINE")
                else: update_status("OFFLINE")
                show_result("Ready for query.")
            elif current_mode == MODE_AI:
                update_status("SEARCHING...")
                # Add API logic here
                utime.sleep(1)
                update_status("READY")
            else:
                try:
                    res = str(eval(input_buf.replace('X','*')))
                    show_result(res)
                except: show_result("ERROR")
        else:
            if len(input_buf) < 25:
                input_buf += k
            update_input(input_buf)
        last_k = k
        utime.sleep(0.15)
    elif not k: last_k = ""
    utime.sleep(0.01)