import time
import machine
from machine import Pin, I2C, I2S
import board_config

class ES7210:
    """MicroPython driver for the ES7210 4-Channel Audio ADC (Microphone Array).
    
    Controls the ES7210 chip over I2C to configure clocks, channels, gain, and format.
    """
    ADDR = 0x40
    
    def __init__(self, i2c):
        self.i2c = i2c
        
    def init(self, sample_rate=16000, bit_width=16):
        """Initializes the ES7210 registers for dual-microphone recording.
        
        Args:
            sample_rate (int): Audio sample rate (e.g. 16000, 44100, 48000).
            bit_width (int): Data bit depth (16 or 24).
            
        Returns:
            bool: True if initialization was successful, False otherwise.
        """
        print("Initializing ES7210 Microphone ADC...")
        try:
            # 1. Reset the chip
            self._write(0x00, 0xFF) # Write all 1s to reset register
            time.sleep_ms(10)
            self._write(0x00, 0x00) # Release reset
            
            # 2. Power management and system configuration
            self._write(0x01, 0x00) # Enable analog power, reference voltage
            self._write(0x11, 0x60) # Enable master clock PLL
            
            # 3. Configure Clock Dividers
            if sample_rate == 16000:
                self._write(0x02, 0x0C) # BCLK divider
                self._write(0x03, 0x10) # LRCK divider
            else: # 44100 / 48000 defaults
                self._write(0x02, 0x04)
                self._write(0x03, 0x08)
                
            # 4. Input Configuration (Enable Mics 1 and 2, power down Mics 3 and 4)
            self._write(0x47, 0x00) # Enable MIC1 / MIC2 analog front-ends
            self._write(0x48, 0xFF) # Power down MIC3 / MIC4 path
            self._write(0x49, 0x0A) # Power up PGA (Programmable Gain Amplifier) 1 and 2
            self._write(0x4A, 0x00) # Power down PGA 3 and 4
            
            # 5. Microphone Gain Settings (+24dB standard)
            # Gain range: 0x00 (0dB) to 0x0F (+45dB) in 3dB steps. 0x08 = +24dB.
            self._write(0x43, 0x08) # Set MIC1 Gain (+24dB)
            self._write(0x44, 0x08) # Set MIC2 Gain (+24dB)
            
            # 6. Set Digital Interface Format (I2S standard format)
            # Bit width: 0x00 = 24-bit, 0x01 = 16-bit, 0x02 = 8-bit, 0x03 = 32-bit
            fmt = 0x01 if bit_width == 16 else 0x00
            self._write(0x13, fmt)  # Set serial output interface format
            self._write(0x14, 0x18) # Enable frame clock / bit clock output
            
            # 7. Unmute ADCs and enable output
            self._write(0x12, 0x00) # Enable ADC digital filters (unmute)
            self._write(0x15, 0x30) # Enable output data pin (SDOUT) active
            
            print("ES7210 initialization complete.")
            return True
        except Exception as e:
            print(f"Failed to initialize ES7210: {e}")
            return False

    def _write(self, reg, val):
        self.i2c.writeto_mem(self.ADDR, reg, bytes([val]))


def record_audio(duration_seconds=5, filename='recording.pcm'):
    """Records raw stereo PCM data from the dual microphones to a file.
    
    Args:
        duration_seconds (int): How long to record in seconds.
        filename (str): Name of output raw PCM file on the device.
    """
    # 1. Use I2C Control Bus from board_config
    i2c = board_config.i2c_bus
    
    # 1a. Setup Master Clock (MCLK) using PWM if configured
    mclk_pwm = None
    if board_config.audio_mclk_pin is not None:
        mclk_pin = Pin(board_config.audio_mclk_pin, Pin.OUT)
        mclk_pwm = machine.PWM(mclk_pin)
        mclk_pwm.freq(board_config.audio_mclk_freq)
        mclk_pwm.duty_u16(32768)
    
    # 2. Configure I2S Receiver
    i2s = I2S(1, 
              sck=Pin(board_config.audio_i2s_sck), 
              ws=Pin(board_config.audio_i2s_ws), 
              sd=Pin(board_config.audio_i2s_rx_sd), 
              mode=I2S.RX, 
              ibuf=16000, 
              rate=16000, 
              bits=16, 
              format=I2S.STEREO)
              
    # 3. Wake up and configure the microphone chip (ES7210 vs ES8311)
    if board_config.audio_mic_codec == "ES7210":
        mic_adc = ES7210(i2c)
        if not mic_adc.init(sample_rate=16000, bit_width=16):
            i2s.deinit()
            if mclk_pwm: mclk_pwm.deinit()
            return False
    else:
        mic_adc = ES8311(i2c)
        if not mic_adc.init(sample_rate=16000):
            i2s.deinit()
            if mclk_pwm: mclk_pwm.deinit()
            return False
        
    print(f"Recording {duration_seconds} seconds of audio...")
    
    # Create reading buffer (reads 100ms chunks: 16000 samples/sec * 2 channels * 2 bytes/sample * 0.1s = 6400 bytes)
    buffer = bytearray(6400)
    
    start_time = time.time()
    total_bytes = 0
    
    try:
        with open(filename, 'wb') as f:
            while (time.time() - start_time) < duration_seconds:
                # Read raw stereo PCM data from I2S
                bytes_read = i2s.readinto(buffer)
                if bytes_read > 0:
                    f.write(buffer[:bytes_read])
                    total_bytes += bytes_read
                    
        print(f"Recording saved successfully to '{filename}' ({total_bytes} bytes).")
        return True
    except Exception as e:
        print(f"Error during recording: {e}")
        return False
    finally:
        # Always release the I2S peripheral resources
        try:
            i2s.deinit()
        except:
            pass
        if mclk_pwm:
            try:
                mclk_pwm.deinit()
            except:
                pass
        print("I2S receiver and MCLK deinitialized.")


class ES8311:
    """MicroPython driver for the ES8311 Audio Codec (Speaker DAC).
    
    Controls the ES8311 chip over I2C to configure clocks, audio format, and volume.
    """
    ADDR = 0x18
    
    def __init__(self, i2c):
        self.i2c = i2c
        
    def init(self, sample_rate=16000):
        """Initializes the ES8311 registers for audio playback.
        
        Args:
            sample_rate (int): Audio sample rate (typically 16000).
            
        Returns:
            bool: True if initialization was successful, False otherwise.
        """
        print("Initializing ES8311 Speaker DAC...")
        try:
            # 1. Reset the chip
            self._write(0x00, 0x1F)
            time.sleep_ms(10)
            self._write(0x00, 0x00)
            time.sleep_ms(10)
            
            # Clock Configuration (16kHz sample rate, MCLK=6.144MHz)
            self._write(0x01, 0x3F) # Enable all clocks, use MCLK pin
            self._write(0x02, 0x48) # pre_div=3, pre_mult=1
            self._write(0x03, 0x10) # fs_mode=0, adc_osr=16
            self._write(0x04, 0x10) # dac_osr=16
            self._write(0x05, 0x00) # adc_div=1, dac_div=1
            self._write(0x06, 0x03) # bclk_div=4 (4-1=3)
            self._write(0x07, 0x00) # lrck_h=0
            self._write(0x08, 0xFF) # lrck_l=255
            
            # Audio Format Configuration (I2S standard format, 16-bit)
            self._write(0x09, 0x0C) # SDP in: 16-bit I2S
            self._write(0x0A, 0x0C) # SDP out: 16-bit I2S
            
            # System / DAC Power Up
            self._write(0x0D, 0x01) # Power up analog circuitry
            self._write(0x0E, 0x02) # Enable analog PGA, enable ADC modulator
            self._write(0x12, 0x00) # Power up DAC
            self._write(0x13, 0x10) # Enable output to HP drive (speaker/hp output)
            self._write(0x1C, 0x6A) # ADC Equalizer bypass
            self._write(0x37, 0x08) # Bypass DAC equalizer
            
            # Set Volume (0xBF = 0dB)
            self._write(0x32, 0xBF)
            
            # Unmute DAC
            self._write(0x31, 0x00)
            
            # Power On
            self._write(0x00, 0x80)
            
            print("ES8311 initialization complete.")
            return True
        except Exception as e:
            print(f"Failed to initialize ES8311: {e}")
            return False
            
    def set_volume(self, val):
        """Sets DAC digital volume (0-100 scale)."""
        # Volume register 0x32 accepts values from 0 (mute) to 255 (+0dB / max volume).
        reg_val = int((val / 100.0) * 255.0)
        reg_val = max(0, min(255, reg_val))
        try:
            self._write(0x32, reg_val)
        except Exception as e:
            print(f"Failed to set volume: {e}")
            
    def _write(self, reg, val):
        self.i2c.writeto_mem(self.ADDR, reg, bytes([val]))


def play_tone(frequency=440, duration_ms=1000, volume=50):
    """Plays a pure sine wave tone on the board speaker.
    
    Args:
        frequency (int): Tone frequency in Hz (e.g. 440 for A4).
        duration_ms (int): Tone duration in milliseconds.
        volume (int): Volume level from 0 to 100.
    """
    import math
    import struct
    
    print(f"Playing tone: {frequency}Hz for {duration_ms}ms (vol={volume})...")
    
    # 1. Setup Master Clock (MCLK) using PWM if configured
    mclk_pwm = None
    if board_config.audio_mclk_pin is not None:
        mclk_pin = Pin(board_config.audio_mclk_pin, Pin.OUT)
        mclk_pwm = machine.PWM(mclk_pin)
        mclk_pwm.freq(board_config.audio_mclk_freq)
        mclk_pwm.duty_u16(32768)
    
    # 2. Use dynamic I2C Control Bus from board_config
    i2c = board_config.i2c_bus
    
    # 3. Initialize the ES8311 DAC
    dac = ES8311(i2c)
    if not dac.init(sample_rate=16000):
        if mclk_pwm: mclk_pwm.deinit()
        return False
        
    dac.set_volume(volume)
    
    # 4. Configure I2S TX
    i2s = I2S(1,
              sck=Pin(board_config.audio_i2s_sck),
              ws=Pin(board_config.audio_i2s_ws),
              sd=Pin(board_config.audio_i2s_tx_sd),
              mode=I2S.TX,
              ibuf=8000,
              rate=16000,
              bits=16,
              format=I2S.STEREO)
              
    # 5. Enable Speaker Amplifier
    on_val = 0 if board_config.audio_amp_active_level == 0 else 1
    off_val = 1 if board_config.audio_amp_active_level == 0 else 0
    amp_pin = Pin(board_config.audio_amp_pin, Pin.OUT, value=on_val)
    
    # 6. Generate sine wave cycle
    # Approximate frequency to make integer number of samples per cycle
    # (avoiding phase clicking)
    N = int(16000 / frequency)
    N = max(4, N) # prevent division by zero or extremely high frequencies
    
    volume_scale = int((volume / 100.0) * 32767)
    
    cycle_data = bytearray()
    for i in range(N):
        val = int(volume_scale * math.sin(2 * math.pi * i / N))
        cycle_data.extend(struct.pack("<hh", val, val)) # Stereo (L/R)
    cycle_bytes = bytes(cycle_data)
    
    # Write to I2S in chunks
    total_samples = int(16000 * duration_ms / 1000)
    total_cycles = int(total_samples / N)
    
    written_cycles = 0
    while written_cycles < total_cycles:
        cycles_to_write = min(total_cycles - written_cycles, 100)
        i2s.write(cycle_bytes * cycles_to_write)
        written_cycles += cycles_to_write
        
    # 7. Clean up
    time.sleep_ms(100) # Let the remaining buffer play out
    amp_pin.value(off_val) # Disable amp
    i2s.deinit()
    if mclk_pwm: mclk_pwm.deinit()
    print("Tone playback complete.")
    return True


def play_wav(filename, volume=50):
    """Plays a standard WAV audio file on the board speaker.
    
    Standard format: 16-bit PCM, 16kHz sample rate (recommended).
    """
    import struct
    
    print(f"Playing WAV: {filename} (vol={volume})...")
    
    try:
        f = open(filename, 'rb')
    except OSError:
        print(f"Error: Cannot open WAV file '{filename}'")
        return False
        
    try:
        # 1. Parse WAV header chunk by chunk
        riff_header = f.read(12)
        if len(riff_header) < 12 or riff_header[0:4] != b'RIFF' or riff_header[8:12] != b'WAVE':
            print("Error: Invalid WAV file format")
            f.close()
            return False
            
        channels = 1
        sample_rate = 16000
        bits = 16
        
        while True:
            chunk_header = f.read(8)
            if len(chunk_header) < 8:
                break
            chunk_id, chunk_size = struct.unpack('<4sI', chunk_header)
            
            if chunk_id == b'fmt ':
                fmt_data = f.read(chunk_size)
                if len(fmt_data) >= 16:
                    audio_format, channels, sample_rate, byte_rate, block_align, bits = struct.unpack('<HHIIHH', fmt_data[:16])
                    if audio_format != 1:
                        print(f"Warning: Non-PCM audio format ({audio_format})")
                else:
                    print("Error: fmt chunk too small")
                    f.close()
                    return False
            elif chunk_id == b'data':
                # Audio data begins immediately after this chunk size
                break
            else:
                # Skip unknown chunk (align to even byte)
                skip_bytes = (chunk_size + 1) & ~1
                f.seek(skip_bytes, 1)
                
        print(f"WAV Info: {sample_rate}Hz, {bits} bits, {'Mono' if channels == 1 else 'Stereo'}")
        
        # 2. Setup Master Clock (MCLK) using PWM if configured
        mclk_pwm = None
        if board_config.audio_mclk_pin is not None:
            mclk_pin = Pin(board_config.audio_mclk_pin, Pin.OUT)
            mclk_pwm = machine.PWM(mclk_pin)
            mclk_pwm.freq(board_config.audio_mclk_freq)
            mclk_pwm.duty_u16(32768)
        
        # 3. Use dynamic I2C Control Bus from board_config
        i2c = board_config.i2c_bus
        
        # 4. Initialize the ES8311 DAC
        dac = ES8311(i2c)
        if not dac.init(sample_rate=16000):
            if mclk_pwm: mclk_pwm.deinit()
            f.close()
            return False
            
        dac.set_volume(volume)
        
        # 5. Configure I2S TX
        i2s_format = I2S.MONO if channels == 1 else I2S.STEREO
        i2s = I2S(1,
                  sck=Pin(board_config.audio_i2s_sck),
                  ws=Pin(board_config.audio_i2s_ws),
                  sd=Pin(board_config.audio_i2s_tx_sd),
                  mode=I2S.TX,
                  ibuf=4096,
                  rate=sample_rate,
                  bits=bits,
                  format=i2s_format)
                  
        # 6. Enable Speaker Amplifier
        on_val = 0 if board_config.audio_amp_active_level == 0 else 1
        off_val = 1 if board_config.audio_amp_active_level == 0 else 0
        amp_pin = Pin(board_config.audio_amp_pin, Pin.OUT, value=on_val)
        
        # 7. Read and stream chunks to I2S
        buf = bytearray(2048)
        while True:
            bytes_read = f.readinto(buf)
            if bytes_read == 0:
                break
            i2s.write(buf[:bytes_read])
            
        # 8. Clean up
        time.sleep_ms(100) # Let the remaining buffer play out
        amp_pin.value(off_val) # Disable amp
        i2s.deinit()
        if mclk_pwm: mclk_pwm.deinit()
        f.close()
        print("WAV playback complete.")
        return True
        
    except Exception as e:
        print(f"Error during WAV playback: {e}")
        try:
            f.close()
        except:
            pass
        return False