import machine
import time
import math
import struct
import array
import gc
import st7789py as st7789
from config import CURRENT_CONFIG
from audio import AudioPlayer, Microphone
from display import Display

# =============================================================================
# 硬件引脚配置 (从 config.py 获取)
# =============================================================================

def main():
    print("\n" + "="*40)
    print("AUDIO & MIC DIAGNOSTIC V5 (Modular & Clean)")
    print("="*40 + "\n")
    
    # 0. 初始化 Boot 按键 (GPIO 0)
    boot_btn = machine.Pin(0, machine.Pin.IN, machine.Pin.PULL_UP)

    # 1. 初始化背光
    # 使用配置中的引脚
    bl_pin = CURRENT_CONFIG.pins.get('bl')
    if bl_pin is not None:
        try:
            bl = machine.Pin(bl_pin, machine.Pin.OUT)
            bl.on()
        except Exception as e: 
            print(f"Backlight error: {e}")

    # 2. 音频测试 (重点排查)
    speaker = AudioPlayer()
    if speaker.i2s:
        # 默认播放马里奥
        speaker.play_mario()
    else:
        print("!!! Speaker initialization failed")

    # 3. 屏幕初始化
    display = Display()

    # 4. 麦克风实时监测
    mic = Microphone()
    print("\n>>> Starting Mic Monitor...")
    
    read_buf = bytearray(4096) 
    
    # UI
    if display.tft:
        display.init_ui()
    
    last_print = time.ticks_ms()
    last_bar_height = 0
    
    # 录音状态变量
    is_recording = False
    recorded_chunks = []
    
    # 调试：打印一次 Boot 键状态
    print(f"Boot Button Initial State: {boot_btn.value()}")

    heartbeat_state = False

    while True:
        try:
            # === 心跳指示器 (右上角) ===
            # 每隔 100ms 翻转一次，证明循环在跑
            if display.tft:
                heartbeat_state = not heartbeat_state
                color = st7789.GREEN if heartbeat_state else st7789.BLACK
                display.tft.fill_rect(230, 0, 10, 10, color)

            # === 按键录音逻辑 (Boot 键按下) ===
            btn_val = boot_btn.value()
            
            # === 按键状态指示器 (左上角) ===
            # 红色表示按下，蓝色表示未按下
            if display.tft:
                btn_color = st7789.RED if btn_val == 0 else st7789.BLUE
                display.tft.fill_rect(0, 0, 10, 10, btn_color)

            if btn_val == 0:
                if not is_recording:
                    print("\n>>> Start Recording (Boot Pressed)...")
                    is_recording = True
                    recorded_chunks = []
                    if display.tft:
                        print(">>> Filling Screen WHITE")
                        display.fill(st7789.WHITE)
                    else:
                        print(">>> Display TFT is None!")
                
                # 录音
                if mic.i2s:
                    num_read = mic.readinto(read_buf)
                    if num_read > 0:
                        try:
                            recorded_chunks.append(bytes(read_buf[:num_read]))
                        except MemoryError:
                            print("Memory Full!")
                continue # 跳过可视化逻辑

            # === 按键释放处理 ===
            elif is_recording:
                print(f"\n>>> Stop Recording. Captured {len(recorded_chunks)} chunks.")
                is_recording = False
                
                if display.tft:
                    display.init_ui()
                
                # 播放录音
                if speaker.i2s and len(recorded_chunks) > 0:
                    print(">>> Playing...")
                    try:
                        cfg = speaker.config
                        # 重新初始化 Speaker (16kHz Mono 16-bit) 以匹配 Mic 数据
                        speaker.i2s.deinit()
                        speaker.i2s = machine.I2S(
                            0,
                            sck=machine.Pin(cfg['bck']),
                            ws=machine.Pin(cfg['ws']),
                            sd=machine.Pin(cfg['sd']),
                            mode=machine.I2S.TX,
                            bits=16,
                            format=machine.I2S.MONO,
                            rate=16000,
                            ibuf=20000,
                        )
                        
                        # 播放数据
                        for chunk in recorded_chunks:
                            # 32-bit Mono -> 16-bit Mono (取高16位)
                            # chunk 是 bytes, 转为 array('h') 方便访问 16-bit word
                            # 32-bit 数据: LowWord, HighWord
                            # 我们需要 HighWord
                            arr = array.array('h', chunk)
                            samples = arr[1::2]
                            speaker.i2s.write(samples)
                            
                    except Exception as e:
                        print(f"Playback error: {e}")
                    
                    # 恢复 Speaker 原始配置
                    if speaker.i2s: speaker.i2s.deinit()
                    speaker._init_audio()
                
                recorded_chunks = []
                gc.collect()

            # === 原有的可视化逻辑 ===
            if mic.i2s:
                num_read = mic.readinto(read_buf)
                if num_read > 0:
                    sum_squares = 0
                    count = num_read // 4 
                    step = 4 
                    samples_checked = 0
                    max_val = 0
                    
                    for i in range(0, count, step):
                        val = struct.unpack_from('<i', read_buf, i*4)[0]
                        # ICS-43434 24-bit 处理
                        val = val >> 8 
                        sum_squares += val * val
                        if abs(val) > max_val: max_val = abs(val)
                        samples_checked += 1
                        
                    if samples_checked > 0:
                        rms = math.sqrt(sum_squares / samples_checked)
                    else:
                        rms = 0
                    
                    if time.ticks_diff(time.ticks_ms(), last_print) > 1000:
                        print(f"Mic Level -> RMS: {int(rms)}, Max: {max_val}")
                        last_print = time.ticks_ms()
                    
                    if display.tft:
                        # 调整缩放比例，让显示更敏感
                        # 你的日志显示安静时 Max ~2000-3000, 说话时 Max ~40000
                        # 我们可以把 Max 40000 映射到满格
                        
                        bar_height = int((max_val / 40000) * 200) 
                        if bar_height > 200: bar_height = 200
                        if bar_height < 0: bar_height = 0
                        
                        last_bar_height = display.update_audio_bar(bar_height, last_bar_height)

            else:
                time.sleep(0.1)
                
        except Exception as e:
            print(f"Loop error: {e}")
            time.sleep(1)

if __name__ == '__main__':
    main()