from pathlib import Path
import json
import os
from dotenv import load_dotenv
import asyncio
from hume import AsyncHumeClient
from hume.expression_measurement.stream import Config
from hume.expression_measurement.stream.socket_client import StreamConnectOptions
from tqdm import tqdm
import glob
from pydub import AudioSegment
import tempfile
from dataclasses import asdict, is_dataclass

def to_serializable(obj):
    if is_dataclass(obj):
        return asdict(obj)
    elif isinstance(obj, dict):
        return {k: to_serializable(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [to_serializable(v) for v in obj]
    elif hasattr(obj, "__dict__"):
        return {k: to_serializable(v) for k, v in obj.__dict__.items() if not k.startswith("_")}
    else:
        return obj

def run(input_dir, output_dir):
    """
    Stage 3: Data analysis and computation using Hume API for voice emotion detection (streaming).
    """
    print(f"Running Stage 3: Reading from {input_dir}, writing to {output_dir}")

    # Load API keys from .env
    load_dotenv()
    HUME_API_KEY = os.getenv("HUME_API_KEY")
    if not HUME_API_KEY:
        raise RuntimeError("HUME_API_KEY not found in .env file.")

    output_dir.mkdir(parents=True, exist_ok=True)

    # Find all mp3 files in input_dir/audio-data recursively
    audio_dir = Path("audio")
    audio_files = [Path(f) for f in glob.glob(str(audio_dir / "**" / "*.mp3"), recursive=True)]
    audio_files = [f for f in audio_files if f.is_file()]

    ### reverse audio_files list
    # audio_files = audio_files[::-1]

    if not audio_files:
        print("No audio files found for analysis.")
        async def noop():
            return
        return noop

    # Prepare output subdirectory for voice emotions
    voice_emotions_dir = output_dir / "voice-emotions-first-five-seconds"
    voice_emotions_dir.mkdir(parents=True, exist_ok=True)

    async def process_file(audio_path):
        parts = audio_path.parts[-3:]  # last three parts of the path
        name = "/".join(parts)
        out_path = voice_emotions_dir / f"{name}.json"

        if out_path.exists():
            return
        # Trim mp3 to first 5 seconds ([0:5000]), or up to end if shorter than 5s
        audio = AudioSegment.from_file(audio_path)
        start_ms = 0
        end_ms = min(len(audio), 5000)
        if end_ms <= 0:
            # If audio is empty, skip processing
            print(f"Audio {audio_path} is empty, skipping.")
            return
        trimmed = audio[start_ms:end_ms]
        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp:
            trimmed.export(tmp.name, format="mp3")
            trimmed_path = tmp.name

        unfinished = True
        while unfinished:
            try:
                client = AsyncHumeClient(api_key=HUME_API_KEY)
                model_config = Config(prosody={})
                stream_options = StreamConnectOptions(config=model_config)
                async with client.expression_measurement.stream.connect(options=stream_options) as socket:
                    result = await socket.send_file(trimmed_path)
                    # Ensure the directory exists
                    os.makedirs(out_path.parent, exist_ok=True)
                    try:
                        serializable_result = to_serializable(result)
                        # No need to offset time fields, since segment starts at 0
                        json.dump(serializable_result, open(out_path, "w"), indent=2, ensure_ascii=False)
                    except Exception:
                        # Fallback: save string representation if not serializable
                        with open(out_path, "w") as f:
                            f.write(str(result))
                os.remove(trimmed_path)
                unfinished = False
            except:
                # Retry on failure
                print(f"Failed to process {audio_path}. Retrying...")
                await asyncio.sleep(10)
        return os.path.basename(audio_path)

    async def main():
        for audio_path in tqdm(audio_files, desc="Processing audio files"):
            await process_file(audio_path)

    return main  # Return the coroutine function for external use

# Add this block at the end of the file
if __name__ == "__main__":
    main_coroutine = run(Path("./audio-data"), Path("./output"))
    import asyncio
    asyncio.run(main_coroutine())