import json
import os
import re
from datetime import datetime
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass

from pydub import AudioSegment
from ttsfm import TTSClient, Voice, AudioFormat

@dataclass
class PodcastSegment:
    """Represents a speech segment of the podcast"""
    speaker: str
    text: str
    emotion: Optional[str] = None
    voice: Optional[str] = None

@dataclass
class PodcastConfig:
    """Configuration for podcast generation"""
    output_dir: str = "podcast_output"
    audio_format: str = "wav"

class PodcastGenerator:
    """Simplified podcast generator for speech segments only"""
    
    def __init__(self, config: Optional[PodcastConfig] = None):
        self.config = config or PodcastConfig()
        self.client = TTSClient()
        self.speaker_voices = {}
        
        # Default voice pool
        self.default_voice_pool = [
            'alloy', 'echo', 'fable', 'nova', 'onyx', 'shimmer'
        ]
        self.voice_pool_index = 0
        
        # Ensure output directory exists
        os.makedirs(self.config.output_dir, exist_ok=True)
    
    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 validate_json_input(self, json_data: Dict) -> bool:
        """Validate the JSON input structure"""
        if "title" not in json_data:
            raise ValueError("Missing required field: title")
        
        if "segments" not in json_data or not json_data["segments"]:
            raise ValueError("Missing or empty segments")
        
        for i, segment in enumerate(json_data["segments"]):
            if not all(key in segment for key in ["speaker", "text"]):
                raise ValueError(f"Segment {i} missing required fields: speaker, text")
        
        return True
    
    def json_to_podcast_segments(self, json_data: Dict) -> List[PodcastSegment]:
        """Convert JSON data to PodcastSegment objects"""
        segments = []
        
        for segment_data in json_data["segments"]:
            segment = PodcastSegment(
                speaker=segment_data["speaker"].upper(),
                text=segment_data["text"],
                emotion=segment_data.get("emotion"),
                voice=None
            )
            segments.append(segment)
        
        return segments
    
    def assign_voices_to_speakers(self, segments: List[PodcastSegment], voice_assignments: Optional[Dict[str, str]] = None):
        """Assign voices to speakers"""
        if voice_assignments:
            self.speaker_voices = {k.upper(): v for k, v in voice_assignments.items()}
        
        # Find all unique speakers
        speakers = set(segment.speaker.upper() for segment in segments)
        
        # 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:
            segment.voice = self.speaker_voices[segment.speaker.upper()]
    
    def generate_speech_segment(self, segment: PodcastSegment) -> str:
        """Generate speech for a single segment and return file path"""
        voice_enum = self.get_voice_enum(segment.voice)
        
        # Try WAV first, fallback to MP3 if needed
        try:
            response = self.client.generate_speech(
                text=segment.text,
                voice=voice_enum,
                response_format=AudioFormat.WAV
            )
            file_extension = "wav"
        except:
            print("WAV format failed, using MP3...")
            response = self.client.generate_speech(
                text=segment.text,
                voice=voice_enum,
                response_format=AudioFormat.MP3
            )
            file_extension = "mp3"
        
        # Save to temporary file
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        speaker_safe = re.sub(r'[^\w\-_]', '_', segment.speaker)
        filename = f"tts_segment_{speaker_safe}_{timestamp}.{file_extension}"
        filepath = os.path.join(self.config.output_dir, filename)
        
        with open(filepath, 'wb') as f:
            f.write(response.audio_data)
        
        return filepath
    
    def combine_audio_files(self, audio_files: List[str]) -> str:
        """Combine audio files using pydub - handles both WAV and MP3"""
        if not audio_files:
            raise ValueError("No audio files to combine")
        
        try:
            # Load first audio file (detect format automatically)
            first_file = audio_files[0]
            if first_file.endswith('.wav'):
                combined_audio = AudioSegment.from_wav(first_file)
            elif first_file.endswith('.mp3'):
                combined_audio = AudioSegment.from_mp3(first_file)
            else:
                combined_audio = AudioSegment.from_file(first_file)
            
            # Add remaining audio files
            for audio_file in audio_files[1:]:
                if audio_file.endswith('.wav'):
                    sound = AudioSegment.from_wav(audio_file)
                elif audio_file.endswith('.mp3'):
                    sound = AudioSegment.from_mp3(audio_file)
                else:
                    sound = AudioSegment.from_file(audio_file)
                combined_audio = combined_audio + sound
            
            # Generate output filename
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_filename = f"combined_podcast_{timestamp}.wav"
            output_path = os.path.join(self.config.output_dir, output_filename)
            
            # Export combined audio as WAV
            combined_audio.export(output_path, format="wav")
            
            return output_path
            
        except Exception as e:
            print(f"pydub failed: {e}")
            print("Falling back to simple file concatenation...")
            return self.simple_combine_audio_files(audio_files)
    
    def simple_combine_audio_files(self, audio_files: List[str]) -> str:
        """Simple fallback method to combine audio files without ffmpeg"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"combined_podcast_{timestamp}.wav"
        output_path = os.path.join(self.config.output_dir, output_filename)
        
        # Simple binary concatenation (works for same format files)
        with open(output_path, 'wb') as output_file:
            for audio_file in audio_files:
                with open(audio_file, 'rb') as input_file:
                    output_file.write(input_file.read())
        
        print(f"Combined {len(audio_files)} files using simple concatenation")
        return output_path
    
    def generate_podcast_from_json(self, json_string: str) -> Tuple[str, Dict]:
        """Generate podcast audio from JSON string"""
        # Parse JSON
        json_data = json.loads(json_string)
        
        # Validate input
        self.validate_json_input(json_data)
        
        # Extract configuration
        title = json_data["title"]
        voice_assignments = json_data.get("voice_assignments", {})
        
        # Convert to podcast segments
        segments = self.json_to_podcast_segments(json_data)
        
        # Assign voices to speakers
        self.assign_voices_to_speakers(segments, voice_assignments)
        
        # Generate audio for each segment
        temp_audio_files = []
        
        for i, segment in enumerate(segments):
            print(f"Processing segment {i+1}/{len(segments)}: {segment.speaker}")
            audio_file = self.generate_speech_segment(segment)
            temp_audio_files.append(audio_file)
        
        # Combine all audio files
        final_audio_path = self.combine_audio_files(temp_audio_files)
        
        # Clean up temporary files
        for temp_file in temp_audio_files:
            try:
                os.remove(temp_file)
            except:
                pass
        
        # Create metadata
        metadata = {
            "title": title,
            "audio_file": final_audio_path,
            "format": "wav",
            "segments_count": len(segments),
            "voice_assignments": self.speaker_voices.copy(),
            "generated_at": datetime.now().isoformat(),
            "file_size_bytes": os.path.getsize(final_audio_path)
        }
        
        return final_audio_path, metadata


def generate_podcast_from_json(json_string: str, output_dir: str = "podcast_output") -> Tuple[str, Dict]:
    """Convenience function to generate podcast from JSON string"""
    config = PodcastConfig(output_dir=output_dir)
    generator = PodcastGenerator(config)
    return generator.generate_podcast_from_json(json_string)


# Example usage
def test_podcast_generation(sample_json=None):
    """Test podcast generation with sample JSON"""
    
    # Use provided sample_json or create a default one
    if sample_json is None:
        sample_json = {
            "title": "Nature's Wonders",
            "episode": "Reflections of Beauty",
            "participants": [
            {
                "name": "Alice",
                "role": "Host"
            },
            {
                "name": "Bob",
                "role": "Co-host"
            }
            ],
            "segments": [
            {
                "speaker": "Alice",
                "text": "Welcome back to another episode of Nature's Wonders! Today, we're diving into the beauty of natural landscapes as we discuss two stunning images."
            },
            {
                "speaker": "Bob",
                "text": "Absolutely, Alice! The first image features a serene river scene surrounded by lush pine trees. It's like a calm oasis in the heart of nature."
            },
            {
                "speaker": "Alice",
                "text": "Right! The reflections in the water are remarkable. Did you know that the clarity of the water helps in capturing those reflections beautifully?"
            },
            {
                "speaker": "Bob",
                "text": "Yes! This phenomenon occurs when the water is still, and there's minimal disturbance. The early morning or late afternoon light also enhances these reflections."
            }
            ]
        }
    
    json_string = json.dumps(sample_json, indent=2)
    
    try:
        audio_path, metadata = generate_podcast_from_json(json_string)
        print(f"Generated podcast: {audio_path}")
        print(f"Metadata: {metadata}")
        return audio_path, metadata
        
    except Exception as e:
        print(f"Error: {str(e)}")
        return None, None

