import time
import sys
import machine
from machine import Pin, SPI, I2C, I2S

try:
    import network
    has_network = True
except ImportError:
    has_network = False
import board_config

class DummyMCP:
    def __init__(self):
        self.override_active = False
        self.active_led_mode = "off"
    def update(self): pass
    def start(self, port=80): pass

class DummyVStream:
    def __init__(self):
        self.active = False
    def update(self): pass
    def start(self): pass

# Import our utility classes
from shtc3_util import SHTC3
from rtc_util import PCF85063
from battery_util import BatteryMonitor
from button_util import BoardButtons
from rgb_led_util import BoardLED
from ble_util import BLEUART
from mcp_server import MCPServer
from video_stream import VideoStreamServer
from audio_util import ES8311
import struct
import urequests

def create_wav_header(data_size):
    # Generates a 44-byte WAV header for 16kHz, 16-bit mono PCM
    riff = b'RIFF'
    file_size = data_size + 36
    wave = b'WAVE'
    fmt = b'fmt '
    chunk_size = 16
    audio_format = 1 # PCM
    channels = 1 # Mono
    sample_rate = 16000
    bits_per_sample = 16
    byte_rate = sample_rate * channels * (bits_per_sample // 8)
    block_align = channels * (bits_per_sample // 8)
    data_label = b'data'
    
    return struct.pack('<4sI4s4sIHHIIHH4sI', 
                       riff, file_size, wave, fmt, chunk_size, 
                       audio_format, channels, sample_rate, byte_rate, 
                       block_align, bits_per_sample, data_label, data_size)

# LED mode options for manual cycling
led_modes = [
    ("Red (Breathing)", lambda led: led.set_color(40, 0, 0), "breath"),
    ("Green (Breathing)", lambda led: led.set_color(0, 40, 0), "breath"),
    ("Blue (Breathing)", lambda led: led.set_color(0, 0, 40), "breath"),
    ("Cyan (Breathing)", lambda led: led.set_color(0, 30, 30), "breath"),
    ("Magenta (Breathing)", lambda led: led.set_color(30, 0, 30), "breath"),
    ("Rainbow Cycle", lambda led: led.set_color(30, 30, 30), "rainbow"),
    ("LED Off", lambda led: led.off(), "off")
]
local_led_mode_idx = 1 # Green breathing

def main():
    global local_led_mode_idx
    import board_config
    
    print("=== Starting MCP Server Main Boot (Type: {}) ===".format(board_config.BOARD_TYPE))
    
    # 1. Use pre-initialized display and buses from board_config
    i2c = board_config.i2c_bus
    spi = board_config.spi_bus
    display = board_config.display_instance
    touch = board_config.touch
    
    # Configure Audio Amp control pin to save power
    if sys.platform != 'rp2' and board_config.audio_amp_pin is not None:
        # Turn OFF amplifier on boot (active-low vs active-high)
        off_val = 1 if board_config.audio_amp_active_level == 0 else 0
        amp_pin = Pin(board_config.audio_amp_pin, Pin.OUT, value=off_val)
    
    # 2. Initialize utility objects
    sensor = None
    rtc_chip = None
    if i2c is not None:
        if board_config.has_sensor:
            try:
                sensor = SHTC3(i2c)
            except Exception as e:
                print("Failed to initialize SHTC3:", e)
        if board_config.has_rtc:
            try:
                rtc_chip = PCF85063(i2c)
            except Exception as e:
                print("Failed to initialize PCF85063:", e)
            
    battery = BatteryMonitor()
    led = BoardLED(board_config.led_pin)
    buttons = None
    if sys.platform != 'rp2':
        buttons = BoardButtons()
        
    ble_name = "ESP32-S3-" + ("RLCD" if board_config.BOARD_TYPE == "WAVESHARE_RLCD" else "Touch")
    ble_uart = BLEUART(name=ble_name)
    
    # Sync system clock from RTC chip
    if rtc_chip:
        try:
            rtc_chip.sync_to_system()
        except Exception as e:
            print("Failed to sync clock:", e)
    
    # 3. Connect to Wi-Fi status check
    ip_addr = "Offline (USB)"
    mcp = DummyMCP()
    vstream = DummyVStream()
    
    if has_network:
        try:
            wlan = network.WLAN(network.STA_IF)
            ip_addr = wlan.ifconfig()[0] if wlan.isconnected() else "Disconnected"
            
            # 4. Start background TCP/UDP Video Streaming Server
            from video_stream import VideoStreamServer
            vstream = VideoStreamServer(display, tcp_port=8081, udp_port=8082)
            vstream.start()
            
            # 4b. Start MCP Server
            from mcp_server import MCPServer
            mcp = MCPServer(display, led, battery, sensor, rtc_chip, ble_uart, vstream=vstream, touch=touch)
            mcp.start(port=80)
        except Exception as ne:
            print("Failed to start network services:", ne)
            
    # Default to Green Breathing
    led_modes[local_led_mode_idx][1](led)
    mcp.active_led_mode = led_modes[local_led_mode_idx][2]
    
    # Last user actions
    last_action_str = "Boot finished."
    force_dashboard_redraw = True
    
    # 5. Register Button Handlers
    def on_key_click():
        global local_led_mode_idx
        local_led_mode_idx = (local_led_mode_idx + 1) % len(led_modes)
        mode_name, color_fn, mode_type = led_modes[local_led_mode_idx]
        color_fn(led)
        mcp.active_led_mode = mode_type
        print(f"Local Button: Cycle LED -> {mode_name}")
        nonlocal last_action_str, force_dashboard_redraw
        last_action_str = f"Local Button: {mode_name}"
        # Disable override when manual button is pressed
        mcp.override_active = False
        force_dashboard_redraw = True

    def on_boot_click():
        print("Local Button: Force Dashboard refresh.")
        nonlocal last_action_str, force_dashboard_redraw
        last_action_str = "Dashboard Refreshed"
        # Disable override when manual button is pressed
        mcp.override_active = False
        force_dashboard_redraw = True

    if buttons:
        buttons.key.on_click(on_key_click)
        buttons.boot.on_click(on_boot_click)
    
    # Loop state
    last_dashboard_update = 0
    dashboard_update_interval_ms = 5000
    last_led_update = 0
    last_wifi_check = 0
    
    print("ESP32 MCP loop running...")
    
    # 6. Main execution loop
    while True:
        now = time.ticks_ms()
        
        # A. Check Wi-Fi status periodically (every 10 seconds) and auto-reconnect
        if has_network and time.ticks_diff(now, last_wifi_check) >= 10000:
            last_wifi_check = now
            if not wlan.isconnected():
                print("Wi-Fi connection lost. Attempting reconnect...")
                last_action_str = "Wi-Fi Disconnected"
                if ip_addr != "Disconnected":
                    ip_addr = "Disconnected"
                    force_dashboard_redraw = True
                try:
                    wlan.connect(wifi_config.WIFI_SSID, wifi_config.WIFI_PASS)
                except Exception as e:
                    print("Wi-Fi reconnect trigger failed:", e)
            elif ip_addr == "Disconnected" or ip_addr == "Offline (USB)":
                ip_addr = wlan.ifconfig()[0]
                print(f"Wi-Fi Connected! IP Address: {ip_addr}")
                last_action_str = f"Wi-Fi Connected: {ip_addr}"
                force_dashboard_redraw = True
        
        # Check touch interaction if available (Push-to-talk Voice Assistant)
        if touch and touch.is_touched():
            print("Touch detected! Starting Hermes Voice Assistant...")
            
            def draw_status_bar(text):
                y_bar = display.height - 25
                display.line(0, y_bar, display.width, y_bar, 1)
                display.fill_rect(0, y_bar + 1, display.width, 24, 0)
                display.text(text, 10, y_bar + 7, 1)
                display.show()
                
            def clear_status_bar():
                display.fill_rect(0, 215, 320, 25, 0)
                display.show()
                
            # State 1: Wait for user to release their finger from the initial touch
            draw_status_bar("Release finger...")
            
            release_start = time.ticks_ms()
            while touch.is_touched():
                if time.ticks_diff(time.ticks_ms(), release_start) > 2000:
                    print("Release timeout occurred")
                    break
                time.sleep_ms(30)
                
            # State 2: Wait for tap-to-start recording
            draw_status_bar("Ready. Tap screen to record...")
            
            tap_started = False
            start_wait = time.ticks_ms()
            while time.ticks_diff(time.ticks_ms(), start_wait) < 10000: # 10s timeout to start recording
                if touch.is_touched():
                    tap_started = True
                    break
                time.sleep_ms(30)
                
            if tap_started:
                # User tapped the screen. Let's record!
                draw_status_bar("Recording: 15s (Tap to stop)")
                
                # 1. Start MCLK PWM and configure ES8311 mic path
                mclk_pin = Pin(4, Pin.OUT)
                mclk_pwm = machine.PWM(mclk_pin)
                mclk_pwm.freq(6144000)
                mclk_pwm.duty_u16(32768)
                
                codec = ES8311(i2c)
                if codec.init(sample_rate=16000):
                    codec.set_volume(80) # Set a balanced volume to prevent speaker saturation
                    try:
                        codec._write(0x14, 0x1A) # Enable analog mic input & PGA
                        codec._write(0x16, 0x01) # Enable +6dB gain boost for mic recording
                        codec._write(0x17, 0xC8) # Set ADC digital volume
                    except Exception as e:
                        print("Failed to set mic gain:", e)
                
                # 2. Configure I2S RX for recording (Mono 16kHz)
                i2s_rx = I2S(1, 
                             sck=Pin(5), 
                             ws=Pin(7), 
                             sd=Pin(6), 
                             mode=I2S.RX, 
                             ibuf=8000, 
                             rate=16000, 
                             bits=16, 
                             format=I2S.MONO)
                
                audio_chunks = []
                total_data_bytes = 0
                buffer = bytearray(1024)
                
                # Track release of the start-tap first
                start_tap_released = False
                rec_start_time = time.ticks_ms()
                max_rec_duration_ms = 15000 # 15 seconds max duration
                max_rec_bytes = 480000 # 15 seconds of mono 16kHz 16-bit PCM (32KB/s)
                last_sec_left = -1
                
                try:
                    while True:
                        # Safety: check duration and byte size
                        elapsed = time.ticks_diff(time.ticks_ms(), rec_start_time)
                        if elapsed >= max_rec_duration_ms or total_data_bytes >= max_rec_bytes:
                            print("Recording stopped: maximum duration/size reached")
                            break
                            
                        # Update countdown
                        sec_left = max(0, 15 - (elapsed // 1000))
                        if sec_left != last_sec_left:
                            last_sec_left = sec_left
                            draw_status_bar(f"Recording: {sec_left}s (Tap to stop)")
                            
                        # Read I2S chunk
                        bytes_read = i2s_rx.readinto(buffer)
                        if bytes_read > 0:
                            audio_chunks.append(bytes(buffer[:bytes_read]))
                            total_data_bytes += bytes_read
                            
                        # Check touch state
                        touched = touch.is_touched()
                        if not start_tap_released:
                            if not touched:
                                start_tap_released = True
                                print("Start-tap released. Listening for stop-tap...")
                        else:
                            if touched:
                                print("Stop-tap detected. Stopping recording...")
                                break
                except Exception as e:
                    print("Error recording:", e)
                finally:
                    i2s_rx.deinit()
                    try:
                        codec._write(0x16, 0x00) # Reset mic gain back to 0dB immediately to prevent playback saturation
                    except:
                        pass
                    
                if total_data_bytes >= 1000:
                    draw_status_bar("Sending & processing...")
                    
                    # Prepend WAV header
                    raw_pcm = b"".join(audio_chunks)
                    wav_data = create_wav_header(len(raw_pcm)) + raw_pcm
                    
                    # Headers
                    headers = {
                        "Authorization": "Bearer mcT1YA1vOr9wXSiHpCYalweEGGZKX-PIfZv2drp8BSg",
                        "Content-Type": "audio/wav",
                        "X-Device-ID": "kitchen-button"
                    }
                    
                    print(f"Posting WAV ({len(wav_data)} bytes) to Hermes API...")
                    try:
                        res = urequests.post("http://192.168.68.126:8642/api/esp32/voice", headers=headers, data=wav_data)
                        print(f"HTTP Status: {res.status_code}")
                        
                        if res.status_code == 200:
                            response_data = res.content
                            print(f"Received response: {len(response_data)} bytes")
                            
                            if response_data[0:4] == b'RIFF' and response_data[8:12] == b'WAVE':
                                idx = response_data.find(b'fmt ')
                                if idx != -1:
                                    fmt_chunk = response_data[idx:idx+24]
                                    audio_format, channels, sample_rate, byte_rate, block_align, bits = struct.unpack('<HHIIHH', fmt_chunk[8:24])
                                    print(f"Playing response WAV: {sample_rate}Hz, {channels}ch, {bits}bit")
                                    
                                    data_idx = response_data.find(b'data')
                                    if data_idx != -1:
                                        audio_start = data_idx + 8
                                        
                                        draw_status_bar("Speaking...")
                                        
                                        # Play response
                                        amp_pin.value(0) # Enable Amp (Active Low)
                                        i2s_format = I2S.MONO if channels == 1 else I2S.STEREO
                                        i2s_tx = I2S(1,
                                                     sck=Pin(5),
                                                     ws=Pin(7),
                                                     sd=Pin(8),
                                                     mode=I2S.TX,
                                                     ibuf=4096,
                                                     rate=sample_rate,
                                                     bits=bits,
                                                     format=i2s_format)
                                        try:
                                            pos = audio_start
                                            chunk_size = 2048
                                            while pos < len(response_data):
                                                chunk = response_data[pos:pos+chunk_size]
                                                i2s_tx.write(chunk)
                                                pos += chunk_size
                                        finally:
                                            time.sleep_ms(150)
                                            amp_pin.value(1) # Disable Amp
                                            i2s_tx.deinit()
                            else:
                                # MP3 warning message on display
                                draw_status_bar("Error: Gateway returned MP3")
                                time.sleep(3)
                        else:
                            draw_status_bar(f"HTTP Error: {res.status_code}")
                            time.sleep(2)
                    except Exception as he:
                        print("Hermes HTTP post failed:", he)
                        draw_status_bar("Connection Error")
                        time.sleep(2)
                
                # Deinit MCLK PWM
                mclk_pwm.deinit()
                
            # Debounce touch release at the very end of wizard
            clear_status_bar()
            release_end = time.ticks_ms()
            while touch.is_touched():
                if time.ticks_diff(time.ticks_ms(), release_end) > 2000:
                    break
                time.sleep_ms(30)
            time.sleep_ms(200)
            
            last_action_str = "Voice query finished"
            force_dashboard_redraw = True
        
        # A. Handle non-blocking MCP client connection updates
        mcp.update()
        
        # B. Handle non-blocking background video stream updates
        vstream.update()
        
        # C. Draw local dashboard (if not overridden by MCP draw text commands or active video stream)
        if not mcp.override_active and not vstream.active:
            if force_dashboard_redraw or time.ticks_diff(now, last_dashboard_update) >= dashboard_update_interval_ms:
                force_dashboard_redraw = False
                last_dashboard_update = now
                
                # Fetch sensor data
                t, h = sensor.read_sensor() if sensor else (None, None)
                t_str = f"{t} C" if t is not None else "Error"
                h_str = f"{h} %" if h is not None else "Error"
                
                # Fetch battery status
                bat_v = battery.read_voltage()
                bat_p = battery.read_percentage()
                bat_str = f"{bat_v:.2f}V ({bat_p}%)" if bat_v is not None else "Error"
                
                # Fetch current time
                dt = rtc_chip.get_datetime() if rtc_chip else None
                time_str = f"{dt[0]:04d}-{dt[1]:02d}-{dt[2]:02d} {dt[4]:02d}:{dt[5]:02d}:{dt[6]:02d}" if dt else "RTC Error"
                
                # Draw standard status dashboard layout
                display.clear(0)
                line_w = display.width - 10
                
                if board_config.BOARD_TYPE == 'WAVESHARE_RLCD' or sys.platform == 'rp2':
                    title_text = "RP2350-TFT MCP SERVER" if sys.platform == 'rp2' else "Waveshare ESP32-S3-RLCD Server"
                    display.text(title_text, 10, 10, 1)
                    display.line(10, 20, line_w, 20, 1)
                    display.text_large("ENVIRONMENT", 15, 30, scale=2, c=1)
                    display.text(f"Temp        : {t_str}", 25, 55, 1)
                    display.text(f"Humid       : {h_str}", 25, 70, 1)
                    display.line(10, 95, line_w, 95, 1)
                    display.text_large("MCP NET CONNECTION", 15, 105, scale=2, c=1)
                    display.text(f"IP Address  : {ip_addr}", 25, 130, 1)
                    display.text(f"Port / Path : 80 /api/mcp", 25, 145, 1)
                    display.text(f"BLE Name    : {ble_name}", 25, 160, 1)
                    display.line(10, 185, line_w, 185, 1)
                    display.text_large("SYSTEM STATUS", 15, 195, scale=2, c=1)
                    display.text(f"Battery     : {bat_str}", 25, 220, 1)
                    display.text(f"Time        : {time_str}", 25, 235, 1)
                    display.line(10, 255, line_w, 255, 1)
                    display.text(f"Status: {last_action_str}", 15, 265, 1)
                else:
                    title_text = "Hosyond ESP32-S3 Server"
                    display.text(title_text, 10, 8, 1)
                    display.line(10, 18, line_w, 18, 1)
                    # Left Column (System & Environment)
                    display.text("SYSTEM & ENV", 10, 28, 1)
                    display.line(10, 38, 150, 38, 1)
                    display.text(f"Temp : {t_str}", 10, 46, 1)
                    display.text(f"Hum  : {h_str}", 10, 58, 1)
                    display.text(f"Bat  : {bat_str}", 10, 70, 1)
                    display.text(f"Time : {time_str[11:19]}", 10, 82, 1)
                    display.text(f"Date : {time_str[0:10]}", 10, 94, 1)
                    # Right Column (Network)
                    display.text("MCP NETWORK", 170, 28, 1)
                    display.line(170, 38, line_w, 38, 1)
                    display.text(f"IP : {ip_addr}", 170, 46, 1)
                    display.text("Port: 80/api/mcp", 170, 58, 1)
                    display.text(f"BLE : {ble_name[-6:]}", 170, 70, 1)
                    # Bottom Status
                    display.line(10, 115, line_w, 115, 1)
                    display.text(f"Status: {last_action_str}", 10, 125, 1)
                display.show()
                
        # D. Update NeoPixel animation smoothly (runs every 50ms)
        if time.ticks_diff(now, last_led_update) >= 50:
            last_led_update = now
            mode = mcp.active_led_mode
            if mode == "breath":
                led.update_breathing(1.5)
            elif mode == "rainbow":
                led.update_rainbow(0.4)
            elif mode == "off":
                led.off()
            
        # Poll rapidly if stream is active, otherwise sleep 50ms to save power
        if vstream.active:
            time.sleep_ms(2)
        else:
            time.sleep_ms(50)

if __name__ == "__main__":
    main()