

import os
import tempfile
import re
import uuid
from datetime import datetime
from dataclasses import dataclass
from typing import List, Dict, Optional, Tuple, Any
import json
import wave
from ttsfm import TTSClient, Voice, AudioFormat
import platform

# Set up FFmpeg path
FFMPEG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ffmpeg")

# Determine OS and set appropriate FFmpeg paths
IS_WINDOWS = platform.system() == "Windows"

if IS_WINDOWS:
    # Windows paths
    FFMPEG_EXECUTABLE = os.path.join(FFMPEG_DIR, "ffmpeg.exe")
    FFPROBE_EXECUTABLE = os.path.join(FFMPEG_DIR, "ffprobe.exe")
else:
    # Linux paths - adjust for the deeper directory structure
    LINUX_FFMPEG_SUBDIR = "ffmpeg-7.0.2-amd64-static"
    LINUX_FFMPEG_PATH = os.path.join(FFMPEG_DIR, LINUX_FFMPEG_SUBDIR)
    
    # Check if the expected directory exists, otherwise fall back to the base directory
    if os.path.exists(LINUX_FFMPEG_PATH):
        FFMPEG_EXECUTABLE = os.path.join(LINUX_FFMPEG_PATH, "ffmpeg")
        FFPROBE_EXECUTABLE = os.path.join(LINUX_FFMPEG_PATH, "ffprobe")
    else:
        # Fallback to base directory
        FFMPEG_EXECUTABLE = os.path.join(FFMPEG_DIR, "ffmpeg")
        FFPROBE_EXECUTABLE = os.path.join(FFMPEG_DIR, "ffprobe")

# Set FFmpeg path for pydub if the executable is found
if os.path.exists(FFMPEG_EXECUTABLE):
  print(f"FFmpeg executable found at: {FFMPEG_EXECUTABLE}")
  try:
    from pydub import AudioSegment
    AudioSegment.converter = FFMPEG_EXECUTABLE
    AudioSegment.ffmpeg = FFMPEG_EXECUTABLE
    AudioSegment.ffprobe = FFPROBE_EXECUTABLE
    print("Successfully configured pydub to use the specified FFmpeg.")
  except ImportError:
    print("pydub not installed, skipping FFmpeg path configuration for it.")
else:
  print(f"Warning: FFmpeg executable not found at {FFMPEG_EXECUTABLE}")
  print("Audio conversion and merging may not work correctly.")

# Directory to store generated audio files
AUDIO_OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "generated_audio")
os.makedirs(AUDIO_OUTPUT_DIR, exist_ok=True)

@dataclass
class ScriptSegment:
    """Represents a parsed segment of the script"""
    segment_type: str  # 'speech', 'sfx', 'music', 'pause'
    speaker: Optional[str] = None
    emotion: Optional[str] = None
    text: Optional[str] = None
    duration: Optional[float] = None  # for SFX/music/pause
    voice: Optional[str] = None  # assigned voice
    description: Optional[str] = None  # for SFX/music descriptions

class MultiVoiceTTSProcessor:
    """
    Enhanced TTS processor that can handle multi-voice scripts with SFX markers
    """
    def __init__(self, api_key=None, output_dir=None, max_chunk_length=4000):
        self.client = TTSClient(api_key=api_key) if api_key else TTSClient()
        self.output_dir = output_dir or AUDIO_OUTPUT_DIR
        self.max_chunk_length = max_chunk_length
        
        # Voice assignments for speakers
        self.speaker_voices = {}
        
        # Default voice mappings (can be customized)
        self.default_voice_pool = [
            'alloy', 'echo', 'fable', 'nova', 'onyx', 'shimmer'
        ]
        self.voice_pool_index = 0
        
        # Audio format settings
        self.default_format = 'mp3'
        
    def parse_script(self, script_text: str) -> List[ScriptSegment]:
        """
        Parse script text with markers like [SPEAKER: NAME], [EMOTION: ...], [SFX: ...]
        
        Args:
            script_text (str): Script with markup
            
        Returns:
            List[ScriptSegment]: Parsed segments
        """
        segments = []
        lines = script_text.strip().split('\n')
        
        current_speaker = None
        current_emotion = None
        
        for line_num, line in enumerate(lines, 1):
            line = line.strip()
            if not line:
                continue
                
            # Extract all markers from the line
            markers = re.findall(r'\[([^:]+):\s*([^\]]+)\]', line)
            
            # Remove markers from the text
            text_content = re.sub(r'\[[^\]]+\]', '', line).strip()
            
            # Process markers
            segment_type = 'speech'
            speaker = current_speaker
            emotion = current_emotion
            duration = None
            description = None
            
            # Track if we found any control markers on this line
            found_control_markers = False
            
            for marker_type, marker_value in markers:
                marker_type = marker_type.upper().strip()
                marker_value = marker_value.strip()
                
                if marker_type == 'SPEAKER':
                    speaker = marker_value.upper()
                    current_speaker = speaker
                    found_control_markers = True
                elif marker_type == 'EMOTION':
                    emotion = marker_value
                    current_emotion = emotion
                    found_control_markers = True
                elif marker_type == 'SFX':
                    segment_type = 'sfx'
                    description = marker_value
                    found_control_markers = True
                    # Try to extract duration
                    duration_match = re.search(r'(\d+(?:\.\d+)?)\s*seconds?', marker_value.lower())
                    if duration_match:
                        duration = float(duration_match.group(1))
                elif marker_type == 'MUSIC':
                    segment_type = 'music'
                    description = marker_value
                    found_control_markers = True
                    # Try to extract duration
                    duration_match = re.search(r'(\d+(?:\.\d+)?)\s*seconds?', marker_value.lower())
                    if duration_match:
                        duration = float(duration_match.group(1))
                elif marker_type == 'PAUSE':
                    segment_type = 'pause'
                    found_control_markers = True
                    try:
                        duration = float(marker_value.replace('s', '').replace('seconds', '').strip())
                    except ValueError:
                        duration = 1.0  # Default pause duration
            
            # Create segment based on what we found
            if segment_type in ['sfx', 'music', 'pause']:
                # Non-speech segments
                segment = ScriptSegment(
                    segment_type=segment_type,
                    speaker=None,  # Non-speech segments don't have speakers
                    emotion=None,
                    text=None,
                    duration=duration,
                    description=description
                )
                segments.append(segment)
                
            elif text_content:
                # Speech segment with actual text
                if not speaker:
                    # Assign a default speaker
                    speaker = "NARRATOR"
                    current_speaker = speaker
                
                segment = ScriptSegment(
                    segment_type='speech',
                    speaker=speaker,
                    emotion=emotion,
                    text=text_content,
                    duration=None,
                    description=None
                )
                segments.append(segment)
            
            elif found_control_markers and not text_content:
                # Line with only control markers (SPEAKER, EMOTION) - don't create segment
                pass
        
        return segments
    
    def assign_voices_to_speakers(self, segments: List[ScriptSegment], voice_assignments: Dict[str, str] = None):
        """
        Assign voices to speakers found in the script
        
        Args:
            segments: Parsed script segments
            voice_assignments: Optional manual voice assignments {speaker: voice}
        """
        if voice_assignments:
            # Convert keys to uppercase for consistency
            self.speaker_voices = {k.upper(): v for k, v in voice_assignments.items()}
        
        # Find all unique speakers
        speakers = set()
        for segment in segments:
            if segment.speaker and segment.segment_type == 'speech':
                speakers.add(segment.speaker.upper())
        
        # Assign voices to unassigned speakers
        for speaker in speakers:
            if speaker not in self.speaker_voices:
                voice = self.default_voice_pool[self.voice_pool_index % len(self.default_voice_pool)]
                self.speaker_voices[speaker] = voice
                self.voice_pool_index += 1
        
        # Update segments with voice assignments
        for segment in segments:
            if segment.speaker and segment.speaker.upper() in self.speaker_voices:
                segment.voice = self.speaker_voices[segment.speaker.upper()]
    
    def get_voice_enum(self, voice_name):
        """Convert voice name string to Voice enum"""
        voice_mapping = {
            'alloy': Voice.ALLOY,
            'ash': Voice.ASH,
            'ballad': Voice.BALLAD,
            'coral': Voice.CORAL,
            'echo': Voice.ECHO,
            'fable': Voice.FABLE,
            'nova': Voice.NOVA,
            'onyx': Voice.ONYX,
            'sage': Voice.SAGE,
            'shimmer': Voice.SHIMMER,
            'verse': Voice.VERSE
        }
        return voice_mapping.get(voice_name.lower(), Voice.ALLOY)
    
    def get_audio_format_enum(self, format_name):
        """Convert format name string to AudioFormat enum"""
        format_mapping = {
            'mp3': AudioFormat.MP3,
            'wav': AudioFormat.WAV,
            'opus': AudioFormat.OPUS,
            'aac': AudioFormat.AAC,
            'flac': AudioFormat.FLAC,
            'pcm': AudioFormat.PCM
        }
        return format_mapping.get(format_name.lower(), AudioFormat.MP3)
    
    def generate_speech_segment(self, segment: ScriptSegment, audio_format='mp3') -> Tuple[bytes, str]:
        """
        Generate speech for a single segment
        
        Args:
            segment: ScriptSegment with speech content
            audio_format: Audio format for output (note: the TTS service may return WAV regardless)
            
        Returns:
            Tuple of (audio_data, temp_filename)
        """
        if not segment.text or not segment.voice:
            raise ValueError(f"Segment must have text and assigned voice.")
        
        # Modify text based on emotion if specified
        text_to_speak = segment.text
        if segment.emotion:
            # You could add emotion-based text modification here
            # For now, we'll just use the original text
            pass
        
        try:
            voice_enum = self.get_voice_enum(segment.voice)
            
            # Request MP3 format from the TTS service
            format_enum = AudioFormat.MP3
            
            # Generate speech
            response = self.client.generate_speech(
                text=text_to_speak,
                voice=voice_enum,
                response_format=format_enum
            )
            
            audio_data = response.audio_data
            
            # Determine the actual format (the service might return WAV even if we request MP3)
            actual_format = 'wav' if audio_data[:4] == b'RIFF' else 'mp3'
            
            # Save to temporary file with the correct extension
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            speaker_safe = re.sub(r'[^\w\-_]', '_', segment.speaker or 'unknown')
            filename = f"tts_segment_{speaker_safe}_{timestamp}.{actual_format}"
            filepath = os.path.join(self.output_dir, filename)
            
            with open(filepath, 'wb') as f:
                f.write(audio_data)
            
            print(f"Generated speech segment in {actual_format} format: {filepath}")
            
            return audio_data, filepath
            
        except Exception as e:
            raise Exception(f"Failed to generate speech for segment: {str(e)}")
    
    def create_silence(self, duration_seconds: float, audio_format='mp3') -> bytes:
        """
        Create silence audio data
        
        Args:
            duration_seconds: Length of silence
            audio_format: Audio format
            
        Returns:
            bytes: Silent audio data
        """
        try:
            import wave
            import struct
            import io
            from io import BytesIO
            
            # Parameters for silence
            sample_rate = 44100
            num_channels = 1
            sample_width = 2  # 16-bit
            num_frames = int(duration_seconds * sample_rate)
            
            # Create silence data (all zeros)
            silence_data = struct.pack('<' + ('h' * num_frames), *([0] * num_frames))
            
            # Create a BytesIO object to hold the WAV data
            output_buffer = BytesIO()
            
            # Write to BytesIO
            with wave.open(output_buffer, 'wb') as wave_file:
                wave_file.setnchannels(num_channels)
                wave_file.setsampwidth(sample_width)
                wave_file.setframerate(sample_rate)
                wave_file.writeframes(silence_data)
                
            # Return the WAV data
            return output_buffer.getvalue()
            
        except Exception as e:
            print(f"Error creating silence: {e}")
            # Return minimal silence data
            return b'\x00'
    
    def create_placeholder_audio(self, description: str, duration: Optional[float], audio_format='mp3') -> bytes:
        """
        Create placeholder audio for SFX/Music (since we can't generate actual SFX)
        This could be silence or a spoken description
        
        Args:
            description: Description of the SFX/Music
            duration: Duration in seconds
            audio_format: Audio format
            
        Returns:
            bytes: Audio data
        """
        # Use a default duration if none is specified
        duration_seconds = duration if duration is not None else 2.0
        
        # Create silence using our wave-based method
        return self.create_silence(duration_seconds, 'wav')
    
    def generate_multi_voice_audio(self, script_text: str, voice_assignments: Dict[str, str] = None, 
                                 audio_format='mp3', segment_pause_ms=200) -> Tuple[bytes, str, List[Dict[str, Any]], Dict[str, str]]:
        """
        Generate multi-voice audio from script text
        
        Args:
            script_text: Script with markup
            voice_assignments: Manual voice assignments {speaker: voice}
            audio_format: Output audio format
            segment_pause_ms: Pause between segments in milliseconds
            
        Returns:
            Tuple of (merged_audio_data, output_filename, segments_info, voice_assignments)
        """
        try:
            # Parse the script
            segments = self.parse_script(script_text)
            
            if not segments:
                raise ValueError("No segments found in script. Check your script format.")
            
            # Assign voices to speakers
            self.assign_voices_to_speakers(segments, voice_assignments)
            
            # Generate audio for each segment
            audio_segments = []
            temp_files = []
            
            # For merging, we'll use WAV format internally
            internal_format = 'wav'
            
            for i, segment in enumerate(segments):
                if segment.segment_type == 'speech' and segment.text:
                    if not segment.voice:
                        continue
                    
                    # Generate speech
                    try:
                        # Generate speech - the service might return WAV even if we request MP3
                        audio_data, temp_file = self.generate_speech_segment(segment, internal_format)
                        audio_segments.append(audio_data)
                        temp_files.append(temp_file)
                    except Exception as e:
                        print(f"Error generating speech for segment {i}: {e}")
                        # Skip this segment or add silence
                        audio_segments.append(self.create_silence(1.0, internal_format))
                    
                elif segment.segment_type in ['sfx', 'music']:
                    # Create placeholder audio
                    audio_data = self.create_placeholder_audio(
                        segment.description, segment.duration, internal_format
                    )
                    audio_segments.append(audio_data)
                    
                elif segment.segment_type == 'pause':
                    # Create silence
                    duration = segment.duration or 1.0
                    audio_data = self.create_silence(duration, internal_format)
                    audio_segments.append(audio_data)
            
            if not audio_segments:
                raise ValueError("No audio segments generated. Check your script and voice assignments.")
            
            # Merge all segments - this will always produce WAV format
            merged_audio = self.merge_audio_segments(audio_segments, internal_format, segment_pause_ms)
            
            # Save merged file - always use .wav extension since our merger only produces WAV
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_filename = f"multi_voice_output_{timestamp}.wav"
            output_filepath = os.path.join(self.output_dir, output_filename)
            
            with open(output_filepath, 'wb') as f:
                f.write(merged_audio)
            
            # Clean up temporary files
            for temp_file in temp_files:
                try:
                    os.remove(temp_file)
                except:
                    pass
            
            # Prepare segment info for API response
            segments_info = []
            for segment in segments:
                segments_info.append({
                    "type": segment.segment_type,
                    "speaker": segment.speaker,
                    "text": segment.text,
                    "emotion": segment.emotion,
                    "description": segment.description,
                    "duration": segment.duration,
                    "voice": segment.voice
                })
            
            return merged_audio, output_filename, segments_info, self.speaker_voices
            
        except Exception as e:
            raise Exception(f"Multi-voice audio generation failed: {str(e)}")
    
    def merge_audio_segments(self, audio_segments: List[bytes], audio_format: str, pause_ms=200) -> bytes:
        """
        Merge multiple audio segments with pauses
        
        Args:
            audio_segments: List of audio data bytes
            audio_format: Audio format
            pause_ms: Pause between segments in milliseconds
            
        Returns:
            bytes: Merged audio data
        """
        # If we only have one segment, return it directly
        if len(audio_segments) <= 1:
            return audio_segments[0] if audio_segments else b''
        
        try:
            import wave
            import io
            from io import BytesIO
            import os
            import tempfile
            
            # Create a temporary directory for our files
            temp_dir = tempfile.mkdtemp()
            
            # Create a list to store file paths
            temp_files = []
            
            # Save each audio segment to a temporary file
            for i, audio_data in enumerate(audio_segments):
                # Check if this is WAV data by looking at the header
                is_wav = audio_data[:4] == b'RIFF'
                
                if not is_wav:
                    # If not WAV, we can't process it with the wave module
                    print(f"Segment {i} is not WAV format. Skipping.")
                    continue
                
                # Save to temporary file
                temp_file = os.path.join(temp_dir, f"segment_{i}.wav")
                with open(temp_file, 'wb') as f:
                    f.write(audio_data)
                temp_files.append(temp_file)
            
            if not temp_files:
                print("No valid WAV files to merge.")
                return audio_segments[0] if audio_segments else b''
            
            # Create a silence file if needed
            silence_file = None
            if pause_ms > 0 and len(temp_files) > 1:
                silence_file = os.path.join(temp_dir, "silence.wav")
                self._create_silence_wav(silence_file, pause_ms / 1000.0)
            
            # Create output file
            output_file = os.path.join(temp_dir, "merged.wav")
            
            # Get parameters from the first file
            with wave.open(temp_files[0], 'rb') as first_wave:
                params = first_wave.getparams()
                
                # Create output wave file
                with wave.open(output_file, 'wb') as output_wave:
                    output_wave.setparams(params)
                    
                    # Write first file
                    first_wave.rewind()
                    output_wave.writeframes(first_wave.readframes(first_wave.getnframes()))
            
            # Add remaining files with silence in between
            for i in range(1, len(temp_files)):
                # Add silence if needed
                if silence_file:
                    self._append_wav_file(output_file, silence_file)
                
                # Add the next segment
                self._append_wav_file(output_file, temp_files[i])
            
            # Read the merged file
            with open(output_file, 'rb') as f:
                merged_data = f.read()
            
            # Clean up temporary files
            for file in temp_files:
                try:
                    os.remove(file)
                except:
                    pass
            
            if silence_file:
                try:
                    os.remove(silence_file)
                except:
                    pass
            
            try:
                os.remove(output_file)
            except:
                pass
            
            try:
                os.rmdir(temp_dir)
            except:
                pass
            
            return merged_data
        
        except Exception as e:
            print(f"Error merging audio with wave module: {e}")
            # Fallback to returning the first segment
            return audio_segments[0]
    
    def _create_silence_wav(self, filename, duration_seconds):
        """Create a simple silence WAV file"""
        try:
            import wave
            import struct
            
            # Parameters for silence
            sample_rate = 44100
            num_channels = 1
            sample_width = 2  # 16-bit
            num_frames = int(duration_seconds * sample_rate)
            
            # Create silence data (all zeros)
            silence_data = struct.pack('<' + ('h' * num_frames), *([0] * num_frames))
            
            # Write to file
            with wave.open(filename, 'wb') as wave_file:
                wave_file.setnchannels(num_channels)
                wave_file.setsampwidth(sample_width)
                wave_file.setframerate(sample_rate)
                wave_file.writeframes(silence_data)
                
            return True
        except Exception as e:
            print(f"Error creating silence WAV: {e}")
            return False

    def _append_wav_file(self, output_file, input_file):
        """Append one WAV file to another"""
        try:
            # Open the output file in append mode
            with wave.open(output_file, 'rb') as output_wav:
                params = output_wav.getparams()
                frames = output_wav.readframes(output_wav.getnframes())
            
            # Open the input file
            with wave.open(input_file, 'rb') as input_wav:
                input_frames = input_wav.readframes(input_wav.getnframes())
            
            # Create a new file with combined frames
            with wave.open(output_file + '.tmp', 'wb') as combined_wav:
                combined_wav.setparams(params)
                combined_wav.writeframes(frames)
                combined_wav.writeframes(input_frames)
            
            # Replace the original file with the new one
            os.replace(output_file + '.tmp', output_file)
            return True
        except Exception as e:
            print(f"Error appending WAV file: {e}")
            return False

# Helper functions for API integration

def generate_from_script(script_text, voice_assignments=None, api_key=None, audio_format='mp3', language='en-US'):
    """
    Generate audio from a script with multiple voices
    
    Args:
        script_text: The script with markup
        voice_assignments: Optional dictionary mapping speakers to voices
        api_key: Optional API key for TTS service
        audio_format: Audio format to generate (mp3, wav, etc.)
        language: Language code (e.g., "en-US", "ne-NP") for TTS
        
    Returns:
        dict: Response with audio file information and segments
    """
    try:
        processor = MultiVoiceTTSProcessor(api_key=api_key)
        
        # Log the language being used
        print(f"Generating multi-voice audio in language: {language}")
        
        audio_data, filepath, segments_info, voice_assignments = processor.generate_multi_voice_audio(
            script_text=script_text,
            voice_assignments=voice_assignments,
            audio_format=audio_format
        )
        
        # Get the filename from the path
        filename = os.path.basename(filepath)
        
        # Include language in the response
        return {
            "status": "success",
            "message": "Multi-voice audio generated successfully",
            "filename": filename,
            "path": filepath,
            "url": f"/api/audio_generation/audio/{filename}",
            "segments": segments_info,
            "voice_assignments": voice_assignments,
            "language": language,
            "duration_seconds": len(segments_info)  # Approximate duration
        }
        
    except Exception as e:
        return {
            "status": "error",
            "message": f"Failed to generate multi-voice audio: {str(e)}"
        }

def get_available_voices():
    """
    Get available voices for multi-voice TTS
    
    Returns:
        dict: Voice options organized by category
    """
    # These are the available voices in the ttsfm library
    voices = {
        "alloy": {"name": "Alloy", "description": "Balanced and neutral voice"},
        "echo": {"name": "Echo", "description": "Bold and authoritative voice"},
        "fable": {"name": "Fable", "description": "Elegant, confident female voice"},
        "nova": {"name": "Nova", "description": "Natural female voice with warmth"},
        "onyx": {"name": "Onyx", "description": "Deep and rich male voice"},
        "sage": {"name": "Sage", "description": "Thoughtful and measured male voice"},
        "shimmer": {"name": "Shimmer", "description": "Bright and uplifting female voice"}
    }
    
    return voices