import os
import asyncio
from openai import AsyncOpenAI
from pydub import AudioSegment
import re
import config
from rate_limiter import RateLimiter

client = AsyncOpenAI(
    api_key=config.OPENAI_API_KEY,
    organization=config.OPENAI_ORGANIZATION,
    default_headers={"OpenAI-Project": config.PROJECT_ID}
)

rate_limiter = RateLimiter(max_calls=config.MAX_CALLS_PER_MINUTE, period=60)

# Set the path to ffmpeg explicitly
AudioSegment.converter = config.FFMPEG_PATH

def split_text(text, max_length=4000):
    """Split the input text into chunks of maximum length."""
    chunks = []
    current_chunk = []
    current_length = 0

    for line in text.split('\n'):
        if current_length + len(line) + 1 <= max_length:
            current_chunk.append(line)
            current_length += len(line) + 1
        else:
            chunks.append('\n'.join(current_chunk))
            current_chunk = [line]
            current_length = len(line)

    if current_chunk:
        chunks.append('\n'.join(current_chunk))

    return chunks

def get_voice_for_marker(marker):
    """Determine the appropriate voice based on the marker."""
    return config.VOICE_OPTIONS.get(marker, config.VOICE_OPTIONS["NARRATOR"])

async def generate_audio_for_text(text, voice, file_path):
    """Generate audio for a given text using the specified voice."""
    try:
        response = await rate_limiter.run(
            client.audio.speech.create(
                model=config.TTS_MODEL,
                voice=voice,
                input=text
            )
        )
        
        with open(file_path, 'wb') as file:
            file.write(response.content)  # Use .content instead of .read()
        
        return AudioSegment.from_mp3(file_path)
    except Exception as e:
        print(f"Error generating audio for text: {e}")
        return None

async def process_dialogue(chunk):
    """Process a chunk of text with dialogue markers."""
    chunk_audio = AudioSegment.empty()
    lines = chunk.split('\n')
    
    for line in lines:
        marker_match = re.match(r'\[(NARRATOR|MALE_\d+|FEMALE_\d+)\]:', line)
        if marker_match:
            marker = marker_match.group(1)
            voice = get_voice_for_marker(marker)
            text = line[len(marker_match.group(0)):].strip()
        else:
            voice = config.VOICE_OPTIONS["NARRATOR"]
            text = line.strip()

        if text:
            temp_file = os.path.join(config.STORAGE_PATH, f"temp_{voice}.mp3")
            line_audio = await generate_audio_for_text(text, voice, temp_file)
            if line_audio:
                chunk_audio += line_audio
            os.remove(temp_file)

    return chunk_audio

async def tts_f(story, use_dialogue=True):
    """Main function to convert text to speech."""
    print(f"Starting TTS process. Story length: {len(story)} characters, Use dialogue: {use_dialogue}")
    
    try:
        chunks = split_text(story)
        audio_files = []

        if not os.path.exists(config.STORAGE_PATH):
            os.makedirs(config.STORAGE_PATH)
            print(f"Created storage directory: {config.STORAGE_PATH}")

        for i, chunk in enumerate(chunks):
            if use_dialogue:
                chunk_audio = await process_dialogue(chunk)
            else:
                voice = config.VOICE_OPTIONS["NARRATOR"]
                temp_file = os.path.join(config.STORAGE_PATH, f"temp_chunk_{i}.mp3")
                chunk_audio = await generate_audio_for_text(chunk, voice, temp_file)
                os.remove(temp_file)

            if chunk_audio:
                chunk_file = os.path.join(config.STORAGE_PATH, f"chunk_{i}.mp3")
                chunk_audio.export(chunk_file, format="mp3")
                audio_files.append(chunk_file)

        # ... rest of the function remains the same

        # Combine all audio files
        combined = AudioSegment.empty()
        for audio_file in audio_files:
            combined += AudioSegment.from_mp3(audio_file)
        
        final_path = os.path.join(config.STORAGE_PATH, config.SPEECH_AUDIO_FILE)
        combined.export(final_path, format="mp3")
        
        # Check if the exported file is valid
        try:
            AudioSegment.from_mp3(final_path)
            print(f"Exported combined audio to: {final_path}")
        except Exception as e:
            print(f"Error: Generated audio file is invalid. {str(e)}")
            return f'An error occurred in TTS.PY: Generated audio file is invalid'
        
        # Clean up individual chunk files
        for file in audio_files:
            os.remove(file)
            print(f"Removed temporary file: {file}")

        return 'TTS generated successfully'
    except Exception as e:
        print(f"Detailed error in tts_f: {str(e)}")
        return f'An error occurred in TTS.PY: {str(e)}'

if __name__ == "__main__":
    test_prompt = """
    [NARRATOR] I am Jack, and the neon-soaked Las Vegas night presses against the windows of my cramped office. The harsh fluorescent lights flicker overhead, casting an eerie glow on the sea of papers sprawled across my desk. The clock on the wall ticks relentlessly, a constant reminder of another late night at the Oasis Resort and Casino.
[NARRATOR] My fingers fly over the keyboard, the gentle clacking a soothing rhythm in the otherwise silent room. Numbers dance across my screen, a complex ballet of debits and credits that only I can truly appreciate. To others, it's just data. To me, it's a story waiting to be unraveled.
[MALE_1] "Still here, Jack? Don't you ever go home?"
    """
    asyncio.run(tts_f(test_prompt, narrator_voice="neutral", use_dialogue=True))