import openai
from openai import OpenAI
import pyaudio
import wave
import threading
from dotenv import load_dotenv
import os
import argparse

# Load environment variables from .env file
load_dotenv()

# Default API config (will be overridden by CLI args if provided)
api_key = os.getenv('OPENAI_API_KEY', '')
api_base_url = os.getenv('OPENAI_API_BASE_URL', '').strip()

# Client will be created after args are parsed
client = None

def get_client(key=None, base_url=None):
    """Get or create the OpenAI client."""
    global client
    effective_key = key or api_key
    effective_base = (base_url or api_base_url).strip()

    if not effective_key:
        raise ValueError("API key is required. Set OPENAI_API_KEY env var or use --api-key")

    if effective_base:
        return OpenAI(api_key=effective_key, base_url=effective_base)
    return OpenAI(api_key=effective_key)

# Model name aliases
MODEL_ALIASES = {
    'cosyvoice2': 'FunAudioLLM/CosyVoice2-0.5B',
    'cosyvoice': 'FunAudioLLM/CosyVoice2-0.5B',
    'tts-1': 'tts-1',
    'tts-1-hd': 'tts-1-hd',
    'gpt-4o-mini-tts': 'gpt-4o-mini-tts',
}

# Map alias to full model name
def resolve_model(model_input):
    """Resolve model alias to full model name."""
    if not model_input:
        return None
    return MODEL_ALIASES.get(model_input.lower(), model_input)

# Check if using SiliconFlow
def is_siliconflow():
    """Check if the API is configured for SiliconFlow."""
    return api_base_url and 'siliconflow' in api_base_url.lower()

# Format voice for CosyVoice2
def format_cosyvoice2_voice(voice, model):
    """Format voice for CosyVoice2: model:voice"""
    if model and 'CosyVoice2' in model and is_siliconflow():
        return f"{model}:{voice}"
    return voice

def list_audio_devices():
    p = pyaudio.PyAudio()
    print("Available audio devices:")
    info = p.get_host_api_info_by_index(0)
    num_devices = info.get('deviceCount')
    # List all available devices, and mark output devices
    for i in range(0, num_devices):
        if p.get_device_info_by_index(i).get('maxOutputChannels') > 0:
            print(f"Device index {i}: {p.get_device_info_by_index(i).get('name')}")
    p.terminate()


def play_saved_audio(file_path, device_index=None):
    # Open the saved audio file
    wf = wave.open(file_path, 'rb')

    print(f"Playing audio to device {device_index}")

    # Setup PyAudio
    p = pyaudio.PyAudio()

    try:
        stream = p.open(format=p.get_format_from_width(wf.getsampwidth()),
                        channels=wf.getnchannels(),
                        rate=wf.getframerate(),
                        output=True,
                        output_device_index=device_index)
        data = wf.readframes(1024)
        while data:
            stream.write(data)
            data = wf.readframes(1024)
    except Exception as e:
        print(f"Error playing audio on device {device_index}: {e}")
    finally:
        stream.stop_stream()
        stream.close()
        wf.close()
        p.terminate()

#Plays to multiple device indexes at the same time
def play_audio_multiplexed(file_paths, device_indices):
    p = pyaudio.PyAudio()
    streams = []
    
    # Open all files and start all streams
    for file_path, device_index in zip(file_paths, device_indices):
        wf = wave.open(file_path, 'rb')
        stream = p.open(format=p.get_format_from_width(wf.getsampwidth()),
                        channels=wf.getnchannels(),
                        rate=wf.getframerate(),
                        output=True,
                        output_device_index=device_index)
        streams.append((stream, wf))
    
    # Play interleaved
    active_streams = len(streams)
    while active_streams > 0:
        for stream, wf in streams:
            data = wf.readframes(1024)
            if data:
                stream.write(data)
            else:
                stream.stop_stream()
                stream.close()
                wf.close()
                active_streams -= 1
    
    p.terminate()
    
def stream_audio_to_virtual_mic(text, voice="fable", model=None, device_index=None, device_index_2=None, api_key=None, api_base=None):
    # Get model from environment variable or use default
    if model is None:
        model = os.getenv('OPENAI_TTS_MODEL', 'tts-1')

    # Format voice for CosyVoice2 if using SiliconFlow
    voice_to_use = format_cosyvoice2_voice(voice, model)
    print(f"Using model: {model}, voice: {voice_to_use}")

    # Get client with potential CLI overrides
    effective_client = get_client(api_key, api_base)

    response = effective_client.audio.speech.create(
        model=model,
        voice=voice_to_use,
        input=text,
        response_format='wav'
    )

    #This can either stream to one device index at a time, or, via multiplexing
    #it can stream to two similtaneously to prevent lag playing in sequence
    if device_index_2 is not None:
        file_path_1 = "output1.wav"
        file_path_2 = "output2.wav"
        response.stream_to_file(file_path_1)
        response.stream_to_file(file_path_2)
        play_audio_multiplexed([file_path_1, file_path_2], [device_index, device_index_2])
    else:
        file_path_1 = "output1.wav"
        response.stream_to_file(file_path_1)
        play_saved_audio(file_path_1, device_index)

    return "";

 

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Text-to-Mic CLI: Convert text to speech and play to virtual microphone',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Examples:
  %(prog)s "Hello world"
  %(prog)s "Hello world" --voice anna --model cosyvoice2
  %(prog)s "Hello world" --voice alex --model FunAudioLLM/CosyVoice2-0.5B --device 8
  %(prog)s --list-devices
  %(prog)s --list-voices

Environment variables (optional, CLI args take precedence):
  OPENAI_API_KEY        - Your API key (required for API TTS)
  OPENAI_API_BASE_URL  - Custom API base URL (e.g., https://api.siliconflow.cn/v1)
  OPENAI_TTS_MODEL      - Default TTS model (default: tts-1)

Model aliases:
  cosyvoice2, cosyvoice  -> FunAudioLLM/CosyVoice2-0.5B
  tts-1                  -> tts-1
  tts-1-hd               -> tts-1-hd
  gpt-4o-mini-tts        -> gpt-4o-mini-tts

SiliconFlow CosyVoice2 voices:
  alex, anna, bella, benjamin, charles, claire, david, diana

OpenAI voices:
  alloy, ash, ballad, coral, echo, fable, onyx, nova, sage, shimmer
'''
    )

    parser.add_argument('text', nargs='?', help='Text to convert to speech')
    parser.add_argument('--voice', '-v', default='fable',
                        help='Voice to use (default: fable). For CosyVoice2: alex, anna, bella, etc.')
    parser.add_argument('--model', '-m', default=None,
                        help='TTS model to use. Use alias (cosyvoice2) or full name (FunAudioLLM/CosyVoice2-0.5B)')
    parser.add_argument('--device', '-d', type=int, default=None,
                        help='Audio device index to play to (use --list-devices to find)')
    parser.add_argument('--device2', type=int, default=None,
                        help='Second audio device index for multiplexed playback')
    parser.add_argument('--list-devices', action='store_true',
                        help='List available audio output devices and exit')
    parser.add_argument('--list-voices', action='store_true',
                        help='List available voices for the configured model and exit')
    parser.add_argument('--list-models', action='store_true',
                        help='List available model aliases and exit')
    parser.add_argument('--api-key', default=None,
                        help='OpenAI API key (or set OPENAI_API_KEY env var)')
    parser.add_argument('--api-base', default=None,
                        help='API base URL (or set OPENAI_API_BASE_URL env var)')

    args = parser.parse_args()

    # Handle list commands
    if args.list_devices:
        list_audio_devices()
        exit(0)

    if args.list_models:
        print("Available model aliases:")
        for alias, full_name in MODEL_ALIASES.items():
            print(f"  {alias:20} -> {full_name}")
        exit(0)

    if args.list_voices:
        # Detect if using SiliconFlow based on API base URL
        if is_siliconflow():
            print("SiliconFlow CosyVoice2 voices:")
            voices = ['alex', 'anna', 'bella', 'benjamin', 'charles', 'claire', 'david', 'diana']
        else:
            print("OpenAI voices:")
            voices = ['alloy', 'ash', 'ballad', 'coral', 'echo', 'fable', 'onyx', 'nova', 'sage', 'shimmer']
        for voice in voices:
            print(f"  {voice}")
        exit(0)

    # Validate text argument
    if not args.text:
        parser.print_help()
        exit(1)

    # Resolve model (alias -> full name)
    model = resolve_model(args.model) if args.model else os.getenv('OPENAI_TTS_MODEL', 'tts-1')
    print(f"Text: {args.text}")
    print(f"Voice: {args.voice}")
    print(f"Model: {model}")

    # Get device index
    device_index = args.device
    device_index_2 = args.device2

    if device_index is None:
        list_audio_devices()
        device_index = int(input("Enter the device index: "))

    stream_audio_to_virtual_mic(args.text, voice=args.voice, model=model,
                                 device_index=device_index, device_index_2=device_index_2,
                                 api_key=args.api_key, api_base=args.api_base)
