import numpy as np
import asyncio
import base64
from io import BytesIO
from goofi.data import Data, DataType
from goofi.node import Node
from goofi.params import FloatParam, StringParam


class VocalExpression(Node):
    def setup(self):
        self.AudioSegment, self.HumeStreamClient, self.BurstConfig, self.ProsodyConfig = import_audio_libs()

    @staticmethod
    def config_input_slots():
        return {"data": DataType.ARRAY}

    @staticmethod
    def config_output_slots():
        return {
            "prosody_label": DataType.STRING,
            "burst_label": DataType.STRING,
            "prosody_score": DataType.ARRAY,
            "burst_score": DataType.ARRAY,
        }

    @staticmethod
    def config_params():
        return {
            "prosody_analysis": {
                "api_key": StringParam("", doc="Hume API key"),
                "emotion_threshold": FloatParam(0.0, 0.0, 1.0, doc="Threshold for filtering emotions"),
            }
        }

    async def decode_emotion_prosody(self, encoded_audio_sample):
        client = self.HumeStreamClient(self.params["prosody_analysis"]["api_key"].value)
        burst_config = self.BurstConfig()
        prosody_config = self.ProsodyConfig()

        prosody_label = ""
        prosody_score = 0.0
        burst_label = ""
        burst_score = 0.0
        emotion_threshold = self.params["prosody_analysis"]["emotion_threshold"].value

        for attempt in range(3):  # Retry up to 3 times
            try:
                async with client.connect([burst_config, prosody_config]) as socket:
                    # Reset connection stream context between samples
                    await socket.reset_stream()

                    # Send binary data
                    result = await socket.send_bytes(encoded_audio_sample)

                    if "prosody" in result:
                        if "warning" in result["prosody"]:
                            print(result["prosody"]["warning"])
                        else:
                            emotions = result["prosody"]["predictions"][0]["emotions"]
                            if emotions:
                                main_emotion = max(emotions, key=lambda e: e["score"])
                                if main_emotion["score"] >= emotion_threshold:
                                    prosody_label = main_emotion["name"]
                                    prosody_score = main_emotion["score"]

                    if "burst" in result:
                        if "warning" in result["burst"]:
                            print(result["burst"]["warning"])
                        else:
                            emotions = result["burst"]["predictions"][0]["emotions"]
                            if emotions:
                                main_emotion = max(emotions, key=lambda e: e["score"])
                                if main_emotion["score"] >= emotion_threshold:
                                    burst_label = main_emotion["name"]
                                    burst_score = main_emotion["score"]

                    return prosody_label, burst_label, prosody_score, burst_score
            except Exception as e:
                print(f"Attempt {attempt + 1} failed: {e}")
                await asyncio.sleep(1)  # Wait before retrying
        raise Exception("Failed to connect to Hume API after multiple attempts")

    def process(self, data: Data):
        if data is None:
            return None

        audio_sample = data.data
        if audio_sample.ndim > 1:
            raise ValueError("Data must be 1D")

        # Convert numpy array to AudioSegment
        audio_segment = self.AudioSegment(
            data=audio_sample.tobytes(),
            sample_width=audio_sample.dtype.itemsize,
            frame_rate=44100,  # Assuming the sample rate is 44100 Hz
            channels=1,  # Assuming the audio is mono
        )

        # Export AudioSegment to WAV format in memory
        audio_io = BytesIO()
        audio_segment.export(audio_io, format="wav")
        audio_io.seek(0)
        encoded_audio_sample = base64.b64encode(audio_io.read()).decode("utf-8")

        prosody_label, burst_label, prosody_score, burst_score = asyncio.run(
            self.decode_emotion_prosody(encoded_audio_sample.encode("utf-8"))
        )

        return {
            "prosody_label": (prosody_label, {}),
            "burst_label": (burst_label, {}),
            "prosody_score": (np.array([prosody_score]), {}),
            "burst_score": (np.array([burst_score]), {}),
        }


def import_audio_libs():
    try:
        from pydub import AudioSegment
    except ImportError:
        raise ImportError("You need to install pydub to use the VocalExpression node: pip install pydub")
    try:
        from hume import HumeStreamClient
        from hume.models.config import BurstConfig, ProsodyConfig
    except ImportError:
        raise ImportError("You need to install hume to use the VocalExpression node: pip install hume")
    return AudioSegment, HumeStreamClient, BurstConfig, ProsodyConfig
