import time
import struct
import machine
from machine import Pin, SPI, I2C, I2S
import urequests
import ili9341
from ft6336u import FT6336U
from audio_util import ES8311

# --- CONFIGURATION ---
HERMES_API_URL = "http://192.168.68.126:8642/api/esp32/voice"
DEVICE_ID = "kitchen-button"

# Placeholder for API Key. Since the key is stored on the Hermes server,
# please copy-paste the token string here.
HERMES_API_KEY = "mcT1YA1vOr9wXSiHpCYalweEGGZKX-PIfZv2drp8BSg"

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)

def main():
    print("--- Starting Hermes Voice Assistant Demo ---")
    
    # 1. Initialize display
    spi = SPI(1, baudrate=40000000, polarity=0, phase=0, sck=Pin(12), mosi=Pin(11), miso=Pin(13))
    display = ili9341.ILI9341(spi, cs=Pin(10), dc=Pin(46), bl=Pin(45), rst=None)
    display.clear(0)
    display.text("Hermes Assistant", 10, 10, 1)
    display.show()
    
    # 2. Initialize touch
    i2c = I2C(0, sda=Pin(16), scl=Pin(15))
    touch = FT6336U(i2c, rst_pin=18, int_pin=17, width=320, height=240, swap_xy=True, invert_x=False, invert_y=True)
    
    # 3. Configure audio codec ES8311
    # We need MCLK pin (GPIO 4) active at 6.144 MHz
    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 not codec.init(sample_rate=16000):
        print("ES8311 init failed!")
        return
        
    codec.set_volume(90)
    
    # Configure microphone registers on ES8311
    try:
        codec._write(0x14, 0x1A) # Enable analog mic input & set PGA gain
        codec._write(0x16, 0x01) # Boost MIC digital gain to +6dB
        codec._write(0x17, 0xC8) # Set ADC digital volume
        print("Microphone registers initialized.")
    except Exception as e:
        print("Failed to write microphone registers:", e)
        
    amp_pin = Pin(1, Pin.OUT, value=1) # start with amp disabled (1 = disabled)
    
    buffer = bytearray(1024)
    
    while True:
        display.clear(0)
        display.text("Hermes Assistant", 10, 10, 1)
        display.line(10, 22, 310, 22, 1)
        display.text("Hold screen & ask", 50, 70, 1)
        display.text("a question...", 50, 90, 1)
        display.text("Status: Idle", 10, 220, 1)
        display.show()
        
        # Wait for touch
        while not touch.is_touched():
            time.sleep_ms(30)
            
        print("Touch detected! Starting recording to RAM...")
        display.clear(0)
        display.text("Hermes Assistant", 10, 10, 1)
        display.line(10, 22, 310, 22, 1)
        display.text("Listening...", 80, 80, 1)
        display.fill_rect(130, 110, 30, 30, 1)
        display.text("Status: Recording", 10, 220, 1)
        display.show()
        
        # 4. Open I2S RX for recording (Mono 16kHz for Hermes STT)
        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
        
        try:
            # Record loop - as long as touch is held
            while touch.is_touched():
                bytes_read = i2s_rx.readinto(buffer)
                if bytes_read > 0:
                    audio_chunks.append(bytes(buffer[:bytes_read]))
                    total_data_bytes += bytes_read
                
            print(f"Touch released! Recorded {total_data_bytes} bytes.")
                
        except Exception as e:
            print("Error recording:", e)
        finally:
            i2s_rx.deinit()
            
        if total_data_bytes < 1000:
            print("Recording too short, ignoring.")
            continue
            
        # 5. Connect and POST to Hermes voice endpoint
        display.clear(0)
        display.text("Hermes Assistant", 10, 10, 1)
        display.line(10, 22, 310, 22, 1)
        display.text("Sending audio...", 50, 80, 1)
        display.text("Waiting for agent...", 50, 100, 1)
        display.text("Status: Processing", 10, 220, 1)
        display.show()
        
        # Join chunks and prepend WAV header
        raw_pcm = b"".join(audio_chunks)
        wav_data = create_wav_header(len(raw_pcm)) + raw_pcm
        
        # Headers
        headers = {
            "Authorization": f"Bearer {HERMES_API_KEY}",
            "Content-Type": "audio/wav",
            "X-Device-ID": DEVICE_ID
        }
        
        print(f"Posting raw WAV ({len(wav_data)} bytes) to {HERMES_API_URL}...")
        try:
            res = urequests.post(HERMES_API_URL, headers=headers, data=wav_data)
            print(f"HTTP Status: {res.status_code}")
            
            if res.status_code == 200:
                response_data = res.content
                print(f"Received {len(response_data)} bytes of audio response.")
                
                # Check response format
                if response_data[0:4] == b'RIFF' and response_data[8:12] == b'WAVE':
                    # Parse WAV header
                    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 WAV: {sample_rate}Hz, {channels}ch, {bits}bit")
                        
                        # Find data chunk
                        data_idx = response_data.find(b'data')
                        if data_idx != -1:
                            audio_start = data_idx + 8
                            
                            display.clear(0)
                            display.text("Hermes Assistant", 10, 10, 1)
                            display.line(10, 22, 310, 22, 1)
                            display.text("Playing response...", 50, 80, 1)
                            display.text("Status: Speaking", 10, 220, 1)
                            display.show()
                            
                            # Enable amp & play
                            amp_pin.value(0) # Active Low Enable
                            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
                                i2s_tx.deinit()
                    else:
                        print("Error: Could not parse WAV format.")
                else:
                    # Check if it is MP3
                    if response_data[0:3] == b'ID3' or (len(response_data) > 2 and response_data[0] == 0xFF and (response_data[1] & 0xE0) == 0xE0):
                        print("Error: Server returned MP3 format. ESP32 requires WAV.")
                        display.clear(0)
                        display.text("Error: Got MP3", 10, 10, 1)
                        display.line(10, 22, 310, 22, 1)
                        display.text("Hermes returned MP3.", 30, 80, 1)
                        display.text("Configure server to", 30, 100, 1)
                        display.text("convert MP3 to WAV.", 30, 120, 1)
                        display.show()
                        time.sleep(4)
                    else:
                        print(f"Error: Unknown response format. Starts with: {response_data[0:16]}")
            else:
                print(f"Request failed with status {res.status_code}: {res.text}")
                display.clear(0)
                display.text("HTTP Error", 10, 10, 1)
                display.line(10, 22, 310, 22, 1)
                display.text(f"Status: {res.status_code}", 50, 80, 1)
                display.show()
                time.sleep(3)
                
        except Exception as e:
            print("Failed to contact Hermes server:", e)
            display.clear(0)
            display.text("Server Error", 10, 10, 1)
            display.line(10, 22, 310, 22, 1)
            display.text("Could not connect", 50, 80, 1)
            display.text("to API gateway.", 50, 100, 1)
            display.show()
            time.sleep(3)
            
        # Debounce touch release
        while touch.is_touched():
            time.sleep_ms(30)
        time.sleep_ms(300)

if __name__ == "__main__":
    main()