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):
        """Generate speech for a single segment and return TTSResponse"""
        voice_enum = self.get_voice_enum(segment.voice)
        
        # Generate speech and return the response object
        response = self.client.generate_speech(
            model="gpt-4o-mini-tts",
            text=segment.text,
            voice=voice_enum,
        )
        
        return response
    
    def combine_audio_responses(self, responses: List, title: str = "podcast") -> str:
        """Combine audio responses using TTSResponse save_to_file method"""
        if not responses:
            raise ValueError("No audio responses to combine")
        
        # Generate output filename
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename_base = f"combined_{title}_{timestamp}"
        output_path = os.path.join(self.config.output_dir, f"{filename_base}.wav")
        
        # If only one response, save it directly
        if len(responses) == 1:
            # Save using the TTSResponse's save_to_file method
            responses[0].save_to_file(output_path)
            print(f"Audio saved to: {output_path}")
        else:
            # Multiple responses - use pydub to merge chunks into single file
            print(f"Processing {len(responses)} audio chunks...")
            
            # Initialize empty audio segment for combining
            combined_audio = None
            saved_chunks = []
            
            # Save ALL chunks in WAV format first
            for i, response in enumerate(responses):
                wav_filename = f"{filename_base}_temp_part{i:03d}.wav"
                wav_path = os.path.join(self.config.output_dir, wav_filename)
                
                try:
                    response.save_to_file(wav_path)
                    if os.path.exists(wav_path) and os.path.getsize(wav_path) > 0:
                        print(f"Saved chunk {i + 1}/{len(responses)} as WAV ({os.path.getsize(wav_path)} bytes)")
                        saved_chunks.append({"path": wav_path, "format": "wav", "index": i})
                    else:
                        print(f"WARNING: Failed to save chunk {i + 1}: file not created or empty")
                except Exception as e:
                    print(f"ERROR: Could not save chunk {i + 1} as WAV: {str(e)}")
            
            # Verify we have all chunks
            if len(saved_chunks) != len(responses):
                print(f"WARNING: Only saved {len(saved_chunks)} out of {len(responses)} chunks!")
                # Continue with what we have but log the issue
            
            # Process the chunks if we have any
            if saved_chunks:
                # Sort chunks by index to ensure correct order
                saved_chunks.sort(key=lambda x: x['index'])
                
                # Initialize with the first chunk
                first_chunk = saved_chunks[0]["path"]
                try:
                    combined_audio = AudioSegment.from_wav(first_chunk)
                    print(f"Started with chunk 1: {len(combined_audio)/1000:.2f}s")
                except Exception as e:
                    print(f"ERROR: Could not load first chunk: {str(e)}")
                    raise Exception(f"Failed to load first audio chunk: {str(e)}")
                
                # Add each subsequent chunk
                for i, chunk_data in enumerate(saved_chunks[1:], 2):
                    try:
                        chunk_path = chunk_data["path"]
                        next_chunk = AudioSegment.from_wav(chunk_path)
                        
                        # Add the chunk to combined audio
                        combined_audio += next_chunk
                        print(f"Combined chunk {i}/{len(saved_chunks)} (duration: {len(next_chunk)/1000:.2f}s)")
                        print(f"Total duration so far: {len(combined_audio)/1000:.2f}s")
                    except Exception as e:
                        print(f"ERROR: Could not add chunk {i}: {str(e)}")
                        # Continue with other chunks instead of failing completely
                
                # Export the final combined audio in WAV format
                try:
                    combined_audio.export(output_path, format="wav")
                    print(f"Combined audio exported to: {output_path} ({len(combined_audio)/1000:.2f}s)")
                    print(f"Final file size: {os.path.getsize(output_path)} bytes")
                except Exception as e:
                    print(f"ERROR: Could not export final audio: {str(e)}")
                    raise Exception(f"Failed to export combined audio: {str(e)}")
                
                # Clean up temporary chunk files
                for chunk_data in saved_chunks:
                    try:
                        if os.path.exists(chunk_data["path"]):
                            os.remove(chunk_data["path"])
                    except OSError as e:
                        print(f"Warning: Could not remove temporary file {chunk_data['path']}: {e}")
                
                print(f"Successfully processed {len(saved_chunks)} audio chunks")
            else:
                raise Exception("Failed to save any audio chunks")
        
        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
        responses = []
        
        for i, segment in enumerate(segments):
            print(f"Processing segment {i+1}/{len(segments)}: {segment.speaker}")
            response = self.generate_speech_segment(segment)
            responses.append(response)
        
        # Combine all audio responses
        final_audio_path = self.combine_audio_responses(responses, title)
        
        # 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):
    """Test podcast generation with sample JSON"""
    
  
    
    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
