#!/usr/bin/env python3
import os
import sys
import json
import base64
import urllib.request
import urllib.error
import subprocess
import wave
import time

# Default ESP32 IP
ESP32_IP = "192.168.68.122"
PORT = 80
URL = f"http://{ESP32_IP}:{PORT}/api/mcp"

def send_mcp_call(method_name, arguments=None):
    """Sends a JSON-RPC tools/call request to the ESP32 server."""
    payload = {
        "jsonrpc": "2.0",
        "method": "tools/call",
        "params": {
            "name": method_name,
            "arguments": arguments or {}
        },
        "id": 1
    }
    req = urllib.request.Request(
        URL,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    try:
        with urllib.request.urlopen(req, timeout=45.0) as res:
            resp_data = res.read().decode("utf-8")
            resp_json = json.loads(resp_data)
            if "error" in resp_json:
                raise RuntimeError(resp_json["error"].get("message", "Unknown board error"))
            
            content_list = resp_json.get("result", {}).get("content", [])
            if len(content_list) > 0 and content_list[0].get("type") == "text":
                return content_list[0].get("text", "")
            return resp_data
    except urllib.error.URLError as e:
        print(f"Network error contacting ESP32 at {URL}: {e.reason}")
        sys.exit(1)

def remote_write_binary(remote_path, file_data):
    """Writes a binary file to the ESP32 by executing a base64-decoder block."""
    b64_str = base64.b64encode(file_data).decode("utf-8")
    code = f"""
import binascii
with open('{remote_path}', 'wb') as f:
    f.write(binascii.a2b_base64('{b64_str}'))
"""
    send_mcp_call("execute_python", {"code": code})

def remote_read_binary(remote_path):
    """Reads a binary file from the ESP32 using a base64-encoder block."""
    code = f"""
import binascii
try:
    with open('{remote_path}', 'rb') as f:
        print(binascii.b2a_base64(f.read()).decode('utf-8').strip())
except Exception as e:
    print("ERROR:", e)
"""
    result = send_mcp_call("execute_python", {"code": code})
    if result.startswith("ERROR:") or "Execution Failed:" in result:
        raise FileNotFoundError(f"Failed to read remote file: {result}")
    
    # Extract the base64 content from the print stdout
    lines = result.split("\n")
    b64_line = ""
    for line in lines:
        if "Execution Succeeded." in line or "Console output:" in line or not line.strip():
            continue
        b64_line = line.strip()
        break
    
    if not b64_line:
         raise ValueError(f"Could not parse base64 audio from board output: {result}")
         
    return base64.b64decode(b64_line)

def convert_pcm_to_wav(pcm_data, wav_path):
    """Converts 16kHz 16-bit stereo raw PCM to standard WAV format."""
    with wave.open(wav_path, "wb") as wav_file:
        wav_file.setnchannels(2)      # Stereo
        wav_file.setsampwidth(2)      # 16-bit (2 bytes)
        wav_file.setframerate(16000)  # 16kHz
        wav_file.writeframes(pcm_data)

def generate_local_speech(text, output_wav):
    """Generates speech WAV file using macOS built-in say and afconvert tools."""
    aiff_path = "temp_speech.aiff"
    try:
        # 1. Synthesize text to AIFF
        subprocess.run(["say", "-o", aiff_path, text], check=True)
        # 2. Convert AIFF to 16kHz 16-bit Mono WAV matching ES8311 speaker driver
        subprocess.run(["afconvert", "-f", "WAVE", "-d", "LEI16@16000", aiff_path, output_wav], check=True)
    finally:
        if os.path.exists(aiff_path):
            os.remove(aiff_path)

def process_command_local(command):
    """Offline keyword parser for testing voice loops without API keys."""
    cmd = command.lower()
    speak_text = ""
    tool_calls = []

    if "led" in cmd or "light" in cmd:
        if "red" in cmd:
            tool_calls.append(("set_led", {"r": 120, "g": 0, "b": 0, "mode": "static"}))
            speak_text = "I have turned the light red."
        elif "green" in cmd:
            tool_calls.append(("set_led", {"r": 0, "g": 120, "b": 0, "mode": "static"}))
            speak_text = "I have turned the light green."
        elif "blue" in cmd:
            tool_calls.append(("set_led", {"r": 0, "g": 0, "b": 120, "mode": "static"}))
            speak_text = "I have turned the light blue."
        elif "rainbow" in cmd:
            tool_calls.append(("set_led", {"r": 30, "g": 30, "b": 30, "mode": "rainbow"}))
            speak_text = "Activating rainbow cycle."
        elif "off" in cmd:
            tool_calls.append(("set_led", {"r": 0, "g": 0, "b": 0, "mode": "off"}))
            speak_text = "Turning off the light."
        else:
            speak_text = "I heard you mention the L E D, but didn't catch the color. You can say red, green, blue, rainbow, or off."
            
    elif "temp" in cmd or "sensor" in cmd or "humid" in cmd:
        # Fetch actual sensor data
        resp = send_mcp_call("get_sensors")
        try:
            data = json.loads(resp)
            t, h = data.get("temperature_c"), data.get("humidity_pct")
            speak_text = f"The current temperature is {t} degrees celsius, and humidity is {h} percent."
        except:
            speak_text = "I failed to read the temperature sensor."
            
    elif "battery" in cmd or "power" in cmd:
        # Fetch actual battery data
        resp = send_mcp_call("get_battery")
        try:
            data = json.loads(resp)
            v, p = data.get("voltage_v"), data.get("percentage_pct")
            speak_text = f"The battery is at {p} percent, with a voltage of {v:.2f} volts."
        except:
            speak_text = "I failed to read the battery status."
            
    elif "beep" in cmd or "tone" in cmd:
        tool_calls.append(("play_tone", {"frequency": 800, "duration_ms": 500, "volume": 50}))
        speak_text = "Beep."
        
    elif "hello" in cmd or "hi there" in cmd:
        speak_text = "Hello! I am your ESP32 voice assistant. How can I help you today?"
        
    else:
        speak_text = f"You said: {command}. I don't have a local action for that command yet."

    return speak_text, tool_calls

def run_openai_assistant(command):
    """Sends the command to OpenAI GPT to decide tools and text responses."""
    try:
        from openai import OpenAI
        client = OpenAI() # Reads OPENAI_API_KEY automatically
    except Exception as e:
        print(f"Failed to initialize OpenAI Client: {e}")
        return None, None

    # Get telemetry data to pass as context
    battery = send_mcp_call("get_battery")
    sensors = send_mcp_call("get_sensors")

    system_prompt = f"""You are a helpful voice assistant running on a Waveshare ESP32-S3 board. 
The user is speaking to you directly. Keep your voice response concise, natural, and conversational (1-2 sentences max), as it will be spoken back via a text-to-speech driver.

You can execute actions on the board by outputting specific tool calls.
Available Actions:
- set_led(r, g, b, mode) - mode can be "static", "breath", "rainbow", or "off"
- play_tone(frequency, duration_ms, volume) - makes a beep
- clear_screen(color) - color is 0 for white, 1 for black
- draw_text(text, x, y, size) - size is 1 or 2

Current Board State:
- Sensors: {sensors}
- Battery: {battery}

To trigger actions, format them in your reply using JSON markers like this:
[TOOL_CALL: {{"name": "set_led", "arguments": {{"r": 120, "g": 0, "b": 0, "mode": "static"}}}} ]

Respond to the user now.
"""
    try:
        completion = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": command}
            ]
        )
        response_text = completion.choices[0].message.content
        
        # Parse tool calls from LLM response
        tool_calls = []
        clean_speech = response_text
        
        import re
        matches = re.finditer(r"\[TOOL_CALL:\s*(\{.*?\})\s*\]", response_text)
        for match in matches:
            try:
                tcall = json.loads(match.group(1))
                tool_calls.append((tcall["name"], tcall.get("arguments", {})))
                clean_speech = clean_speech.replace(match.group(0), "")
            except Exception as ex:
                print("Failed to parse tool call JSON:", ex)
                
        return clean_speech.strip(), tool_calls
    except Exception as e:
        print("OpenAI LLM request failed:", e)
        return None, None

def transcribe_audio_whisper(wav_path):
    """Transcribes WAV file using OpenAI Whisper API (if key is set) or local model (fallback)."""
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key:
        print("Sending audio to OpenAI Whisper API...")
        try:
            from openai import OpenAI
            client = OpenAI()
            with open(wav_path, "rb") as audio_file:
                transcription = client.audio.transcriptions.create(
                    model="whisper-1",
                    file=audio_file
                )
            return transcription.text
        except Exception as e:
            print("OpenAI Whisper API failed:", e)
            
    # Local fallback
    print("Attempting local transcription via whisper package...")
    try:
        import whisper
        # Load small/tiny model for speed
        model = whisper.load_model("tiny")
        result = model.transcribe(wav_path)
        return result["text"]
    except Exception as e:
        print(f"Local Whisper transcription failed: {e}")
        return None

def main():
    print("====================================================")
    print(" ESP32-S3 Voice Assistant Host Coordinator Started")
    print(f" Connecting to Board: {URL}")
    print("====================================================")

    # 1. Quick Connection Verification
    try:
        battery = send_mcp_call("get_battery")
        print(f"Connected to ESP32! Battery level: {json.loads(battery).get('percentage_pct')}%")
    except Exception as e:
        print(f"Connection failed: {e}")
        print("Please verify the board is powered on, connected to the same Wi-Fi, and the IP is correct.")
        return

    while True:
        print("\nReady for command. Choose an option:")
        print(" [Enter] Record a voice command (4 seconds) on the board")
        print(" [t]     Type a text command manually")
        print(" [q]     Quit")
        
        choice = input("Select: ").strip().lower()
        if choice == "q":
            break
        
        command_text = ""
        
        if choice == "t":
            command_text = input("Type command: ").strip()
            if not command_text:
                continue
        else:
            # Voice recording path
            duration = 4
            pcm_filename = "recording.pcm"
            wav_filename = "temp_recording.wav"
            
            print(f"\nRecording voice command for {duration} seconds on the board...")
            print(">>> SPEAK NOW! <<<")
            
            try:
                # Trigger recording on the board (blocks until finished)
                rec_msg = send_mcp_call("record_voice", {"duration_sec": duration, "filename": pcm_filename})
                print(rec_msg)
                
                print("Downloading recording from board over Wi-Fi...")
                pcm_data = remote_read_binary(pcm_filename)
                
                print("Converting PCM to WAV...")
                convert_pcm_to_wav(pcm_data, wav_filename)
                
                print("Transcribing voice...")
                command_text = transcribe_audio_whisper(wav_filename)
                
                if command_text:
                    print(f"Transcribed Text: \"{command_text}\"")
                else:
                    print("Could not transcribe any speech.")
                    continue
            except Exception as err:
                print(f"Error capturing speech: {err}")
                continue
            finally:
                # Clean up local temp recording file
                if os.path.exists(wav_filename):
                    os.remove(wav_filename)

        # 2. Process Command (OpenAI or Local fallbacks)
        speak_text = ""
        tool_calls = []
        
        use_openai = bool(os.environ.get("OPENAI_API_KEY"))
        if use_openai:
            print("Querying OpenAI assistant...")
            speak_text, tool_calls = run_openai_assistant(command_text)
            
        if not speak_text:
            if use_openai:
                print("OpenAI processing failed. Falling back to local offline parser.")
            else:
                print("No OPENAI_API_KEY found. Using offline keyword parser.")
            speak_text, tool_calls = process_command_local(command_text)

        # 3. Execute Tool Actions on the Board
        if tool_calls:
            for tool_name, tool_args in tool_calls:
                print(f"Executing board action: {tool_name}({tool_args})")
                try:
                    res = send_mcp_call(tool_name, tool_args)
                    print(f"Result: {res}")
                except Exception as ex:
                    print(f"Failed to execute tool {tool_name}: {ex}")

        # 4. Generate Voice Response & Play it on the Board
        if speak_text:
            print(f"Assistant speech response: \"{speak_text}\"")
            local_wav = "temp_response.wav"
            remote_wav = "response.wav"
            
            try:
                print("Synthesizing speech WAV file locally...")
                generate_local_speech(speak_text, local_wav)
                
                print("Uploading speech WAV to board...")
                with open(local_wav, "rb") as f:
                    wav_bytes = f.read()
                remote_write_binary(remote_wav, wav_bytes)
                
                print("Playing speech on board speaker...")
                play_res = send_mcp_call("play_audio", {"filename": remote_wav, "volume": 60})
                print(play_res)
                
            except Exception as e:
                print(f"Failed to play voice response: {e}")
            finally:
                if os.path.exists(local_wav):
                    os.remove(local_wav)

if __name__ == "__main__":
    main()