import machine
import time
import struct
import gc
import network
import st7789py as st7789
from config import CURRENT_CONFIG
from audio import AudioPlayer, Microphone
from display import Display
from websocket_client import WebSocketClient
import uselect

WIFI_SSID = "Tangledup-AI"
WIFI_PASS = "djt12345678"
SERVER_IP = "6.6.6.88"
SERVER_PORT = 8000
SERVER_URL = f"ws://{SERVER_IP}:{SERVER_PORT}/ws/audio"


def connect_wifi(max_retries=3):
    wlan = network.WLAN(network.STA_IF)
    
    try:
        wlan.active(False)
        time.sleep(1)
        wlan.active(True)
        time.sleep(1)
    except Exception as e:
        print(f"WiFi init error: {e}")
        return False
    
    for attempt in range(max_retries):
        try:
            if wlan.isconnected():
                print('WiFi connected')
                return True
            
            print(f'Connecting to WiFi {WIFI_SSID}...')
            wlan.connect(WIFI_SSID, WIFI_PASS)
            
            start_time = time.ticks_ms()
            while not wlan.isconnected():
                if time.ticks_diff(time.ticks_ms(), start_time) > 20000:
                    print("WiFi timeout!")
                    break
                time.sleep(0.5)
            
            if wlan.isconnected():
                print('WiFi connected!')
                return True
            
            if attempt < max_retries - 1:
                wlan.disconnect()
                time.sleep(2)
                
        except Exception as e:
            print(f"WiFi error: {e}")
            if attempt < max_retries - 1:
                time.sleep(3)
    
    print("WiFi connection failed!")
    return False


def print_asr(text, display=None):
    print(f"ASR: {text}")
    if display and display.tft:
        display.fill_rect(0, 40, 240, 160, st7789.BLACK)
        display.text(text, 0, 40, st7789.WHITE)


def main():
    print("\n=== ESP32 Audio ASR ===\n")
    
    boot_btn = machine.Pin(0, machine.Pin.IN, machine.Pin.PULL_UP)
    
    bl_pin = CURRENT_CONFIG.pins.get('bl')
    if bl_pin:
        try:
            bl = machine.Pin(bl_pin, machine.Pin.OUT)
            bl.on()
        except:
            pass
    
    speaker = AudioPlayer()
    mic = Microphone()
    display = Display()
    
    if display.tft:
        display.init_ui()
    
    is_recording = False
    ws = None
    
    def connect_ws():
        nonlocal ws
        try:
            if ws:
                ws.close()
        except:
            pass
        ws = None
        
        try:
            print(f"Connecting to {SERVER_URL}")
            ws = WebSocketClient(SERVER_URL)
            print("WebSocket connected!")
            if display:
                display.set_ws(ws)
            return True
        except Exception as e:
            print(f"WS connection failed: {e}")
            return False
    
    if connect_wifi():
        connect_ws()
    else:
        print("Running in offline mode")
    
    read_buf = bytearray(4096)
    
    while True:
        try:
            btn_val = boot_btn.value()
            
            if btn_val == 0:
                if not is_recording:
                    print(">>> Recording...")
                    is_recording = True
                    if display.tft:
                        display.fill(st7789.WHITE)
                    
                    if ws is None or not ws.is_connected():
                        connect_ws()
                    
                    if ws and ws.is_connected():
                        try:
                            ws.send("START_RECORDING")
                        except:
                            ws = None
                
                if mic.i2s:
                    num_read = mic.readinto(read_buf)
                    if num_read > 0:
                        if ws and ws.is_connected():
                            try:
                                ws.send(read_buf[:num_read], opcode=2)
                                
                                poller = uselect.poll()
                                poller.register(ws.sock, uselect.POLLIN)
                                events = poller.poll(0)
                                if events:
                                    msg = ws.recv()
                                    if isinstance(msg, str) and msg.startswith("ASR:"):
                                        print_asr(msg[4:], display)
                                        
                            except:
                                ws = None
                
                continue
            
            elif is_recording:
                print(">>> Stop")
                is_recording = False
                
                if display.tft:
                    display.init_ui()
                
                if ws:
                    try:
                        ws.send("STOP_RECORDING")
                        
                        t_wait = time.ticks_add(time.ticks_ms(), 500)
                        while time.ticks_diff(t_wait, time.ticks_ms()) > 0:
                            poller = uselect.poll()
                            poller.register(ws.sock, uselect.POLLIN)
                            events = poller.poll(100)
                            if events:
                                msg = ws.recv()
                                if isinstance(msg, str) and msg.startswith("ASR:"):
                                    print_asr(msg[4:], display)
                    except:
                        ws = None
                
                gc.collect()
            
            time.sleep(0.01)
            
        except Exception as e:
            print(f"Error: {e}")
            time.sleep(1)


if __name__ == '__main__':
    main()