import json
import os
import re
from datetime import datetime
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass
import asyncio
import aiohttp
import aiofiles

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

try:
    from g4f.client import AsyncClient
except ImportError:
    print("Warning: g4f.client could not be imported. Image generation will be disabled.")
    AsyncClient = None

@dataclass
class PodcastSegment:
    """Podcast ko euta speech segment represent garne class"""
    speaker: str
    text: str
    emotion: Optional[str] = None
    voice: Optional[str] = None
    image_gen_prompt: Optional[str] = None
    image_path: Optional[str] = None
    audio_path: Optional[str] = None
    index: int = 0

@dataclass
class PodcastConfig:
    """Podcast generate garne configuration settings"""
    output_dir: str = "podcast_output"
    audio_dir: str = "podcast_output/audio"
    image_dir: str = "podcast_output/images"
    temp_audio_dir: str = "podcast_output/temp_audio"
    temp_image_dir: str = "podcast_output/temp_images"
    audio_format: str = "wav"
    enable_images: bool = True
    job_id: Optional[str] = None  # Job ID optional ho

class PodcastGenerator:
    """Audio ra images dubai generate garne podcast generator"""
    
    def __init__(self, config: Optional[PodcastConfig] = None):
        self.config = config or PodcastConfig()
        self.client = TTSClient()
        self.speaker_voices = {}
        
        # Available cha bhane image generation client initialize garne
        self.image_client = AsyncClient() if AsyncClient else None
        
        # Default voice pool - customize garna sakincha
        self.default_voice_pool = [
            'alloy', 'echo', 'fable', 'nova', 'onyx', 'shimmer'
        ]
        self.voice_pool_index = 0
        
        # Output directories haru exist garchan ki check garne
        os.makedirs(self.config.output_dir, exist_ok=True)
        os.makedirs(self.config.audio_dir, exist_ok=True)
        os.makedirs(self.config.image_dir, exist_ok=True)
        os.makedirs(self.config.temp_audio_dir, exist_ok=True)
        os.makedirs(self.config.temp_image_dir, exist_ok=True)
    
    def get_voice_enum(self, voice_name):
        """Voice name lai Voice enum ma convert garne"""
        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:
        """JSON input ko structure validate garne"""
        print("JSON Data: ", json_data) 
        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]:
        """JSON data lai PodcastSegment objects ma convert garne"""
        segments = []
        
        for i, segment_data in enumerate(json_data["segments"]):
            segment = PodcastSegment(
                speaker=segment_data["speaker"].upper(),
                text=segment_data["text"],
                emotion=segment_data.get("emotion"),
                voice=None,
                image_gen_prompt=segment_data.get("image_gen_prompt"),
                index=i
            )
            segments.append(segment)
        
        return segments
    
    def assign_voices_to_speakers(self, segments: List[PodcastSegment], voice_assignments: Optional[Dict[str, str]] = None):
        """Speakers haru lai voices assign garne"""
        if voice_assignments:
            self.speaker_voices = {k.upper(): v for k, v in voice_assignments.items()}
        
        # Unique speakers haru find garne
        speakers = set(segment.speaker.upper() for segment in segments)
        
        # Unassigned speakers lai voice assign garne
        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
        
        # Voice assignments le segments update garne
        for segment in segments:
            segment.voice = self.speaker_voices[segment.speaker.upper()]
    
    def generate_speech_segment(self, segment: PodcastSegment, index: int) -> str:
        """Euta segment ko speech generate garne ra file path return garne"""
        voice_enum = self.get_voice_enum(segment.voice)
        
        # WAV format try garne, problem bhayo bhane MP3 ma fallback garne
        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"
        
        # Segment index sanga file ma save garne
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        speaker_safe = re.sub(r'[^\w\-_]', '_', segment.speaker)
        
        # Job ID available cha bhane filename ma include garne
        if self.config.job_id:
            filename = f"segment_{index+1:03d}_{speaker_safe}_{self.config.job_id}_{timestamp}.{file_extension}"
        else:
            filename = f"segment_{index+1:03d}_{speaker_safe}_{timestamp}.{file_extension}"
            
        filepath = os.path.join(self.config.temp_audio_dir, filename)
        
        with open(filepath, 'wb') as f:
            f.write(response.audio_data)
        
        return filepath
    
    async def generate_image_for_segment(self, segment: PodcastSegment, index: int) -> Optional[str]:
        """Euta segment ko image generate garne ra file path return garne"""
        if not self.config.enable_images or not self.image_client or not segment.image_gen_prompt:
            return None
            
        try:
            print(f"Generating image for segment {index+1}...")
            
            # Image generate garne
            response = await self.image_client.images.generate(
                prompt=segment.image_gen_prompt+ "Focus mainly on visual elements and design . Make it more attractive and engaging. Add less text in the image as possible.",
                model="midjourney",
                width=1920,
                height=1080,
                response_format="url"
            )
            
            if response and response.data and response.data[0].url:
                image_url = response.data[0].url
                print(f"Generated image URL for segment {index+1}")
                
                # Image download garne
                speaker_safe = re.sub(r'[^\w\-_]', '_', segment.speaker)
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                
                # Job ID available cha bhane filename ma include garne
                if self.config.job_id:
                    filename = f"segment_{index+1:03d}_{speaker_safe}_{self.config.job_id}_{timestamp}.jpg"
                else:
                    filename = f"segment_{index+1:03d}_{speaker_safe}_{timestamp}.jpg"
                    
                filepath = os.path.join(self.config.temp_image_dir, filename)
                
                if await self.download_image(image_url, filepath):
                    print(f"Saved image for segment {index+1} to {filepath}")
                    return filepath
            
            return None
            
        except Exception as e:
            print(f"Error generating image for segment {index+1}: {e}")
            return None
    
    async def download_image(self, url: str, filepath: str) -> bool:
        """Image ko URL lai download garne"""
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(url) as response:
                    if response.status == 200:
                        content = await response.read()
                        async with aiofiles.open(filepath, 'wb') as f:
                            await f.write(content)
                        return True
                    else:
                        print(f"Failed to download image: HTTP {response.status}")
                        return False
        except Exception as e:
            print(f"Error downloading image: {e}")
            return False
    
    def combine_audio_files(self, audio_files: List[str]) -> str:
        """Audio files lai pydub ma combine garne - WAV ra MP3 support garne"""
        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
            
            # Output filename generate garne
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            
            # Job ID available cha bhane filename ma include garne
            if self.config.job_id:
                output_filename = f"combined_podcast_{self.config.job_id}_{timestamp}.wav"
            else:
                output_filename = f"combined_podcast_{timestamp}.wav"
                
            output_path = os.path.join(self.config.output_dir, output_filename)
            
            # Combined audio as WAV export garne
            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
    
    async def process_segment(self, segment: PodcastSegment, index: int) -> Tuple[str, Optional[str]]:
        """Euta segment ko audio ra image generate garne"""
        # Image generation task start garne, agadi image prompt available cha bhane
        image_task = None
        if segment.image_gen_prompt and self.config.enable_images and self.image_client:
            image_task = asyncio.create_task(self.generate_image_for_segment(segment, index))
        
        # Audio generate garne (synchronously)
        print(f"Processing audio for segment {index+1}/{len(self.segments)}: {segment.speaker}")
        audio_path = self.generate_speech_segment(segment, index)
        segment.audio_path = audio_path
        
        # Image generation task ko wait garne, agadi start gareko
        image_path = None
        if image_task:
            image_path = await image_task
            segment.image_path = image_path
            
        return audio_path, image_path
    
    async def generate_podcast_from_json_async(self, json_string: str) -> Tuple[str, Dict, List[Dict]]:
        """JSON string lai podcast audio ra images generate garne"""
        # JSON parse garne
        json_data = json.loads(json_string)
        
        # Input validate garne
        self.validate_json_input(json_data)
        
        # Configuration extract garne
        title = json_data["title"]
        voice_assignments = json_data.get("voice_assignments", {})
        
        # JSON data lai podcast segments ma convert garne
        self.segments = self.json_to_podcast_segments(json_data)
        
        # Speakers lai voices assign garne
        self.assign_voices_to_speakers(self.segments, voice_assignments)
        
        # Each segment ko process garne (audio ra image)
        temp_audio_files = []
        image_files = []
        segment_results = []
        
        # Segments concurrently process garne
        for i, segment in enumerate(self.segments):
            # Each segment ko one by one process garne, but audio ra image concurrently
            audio_path, image_path = await self.process_segment(segment, i)
            temp_audio_files.append(audio_path)
            
            # Add to the ordered results list
            segment_result = {
                "index": i,
                "speaker": segment.speaker,
                "text": segment.text,
                "emotion": segment.emotion,
                "voice": segment.voice,
                "audio_path": audio_path,
                "image_path": image_path,
                "image_prompt": segment.image_gen_prompt
            }
            segment_results.append(segment_result)
            
            if image_path:
                image_files.append(image_path)
        
        # All audio files combine garne
        final_audio_path = self.combine_audio_files(temp_audio_files)
        
        # Metadata generate garne
        metadata = {
            "title": title,
            "combined_audio_file": final_audio_path,
            "format": "wav",
            "segments_count": len(self.segments),
            "voice_assignments": self.speaker_voices.copy(),
            "image_files": image_files,
            "audio_files": temp_audio_files,
            "generated_at": datetime.now().isoformat(),
            "file_size_bytes": os.path.getsize(final_audio_path),
            "job_id": self.config.job_id
        }
        
        return final_audio_path, metadata, segment_results
    
    def cleanup_temp_files(self, audio_files: Optional[List[str]] = None, image_files: Optional[List[str]] = None):
        """Temporary files cleanup garne after podcast generation"""
        # Temp audio files cleanup garne
        if audio_files:
            for audio_file in audio_files:
                try:
                    if os.path.exists(audio_file):
                        os.remove(audio_file)
                        print(f"Removed temp audio file: {audio_file}")
                except Exception as e:
                    print(f"Failed to remove audio file {audio_file}: {e}")
        
        # Temp image files cleanup garne
        if image_files:
            for image_file in image_files:
                try:
                    if image_file and os.path.exists(image_file):
                        os.remove(image_file)
                        print(f"Removed temp image file: {image_file}")
                except Exception as e:
                    print(f"Failed to remove image file {image_file}: {e}")

    def cleanup_job_files(self):
        """All temporary files cleanup garne based on job_id"""
        if not self.config.job_id:
            print("No job ID specified for cleanup")
            return
            
        job_id = self.config.job_id
        print(f"Cleaning up files for job {job_id}...")
        
        # Find and remove audio files with this job_id
        for temp_dir in [self.config.temp_audio_dir, self.config.temp_image_dir]:
            if os.path.exists(temp_dir):
                print(f"Checking directory: {temp_dir}")
                files_before = len(os.listdir(temp_dir))
                files_removed = 0
                
                for filename in os.listdir(temp_dir):
                    if job_id in filename:
                        try:
                            file_path = os.path.join(temp_dir, filename)
                            if os.path.isfile(file_path):
                                os.remove(file_path)
                                files_removed += 1
                                print(f"Removed file: {filename}")
                        except Exception as e:
                            print(f"Failed to remove {filename}: {e}")
                
                print(f"Cleanup stats for {temp_dir}: {files_removed} files removed out of {files_before} total files")
        
        print(f"Completed cleanup for job {job_id}")
    
    def generate_podcast_from_json(self, json_string: str) -> Tuple[str, Dict, List[Dict]]:
        """Synchronous wrapper for podcast audio generation from JSON string"""
        return asyncio.run(self.generate_podcast_from_json_async(json_string))


async def generate_podcast_from_json_async(json_string: str, output_dir: str = "podcast_output") -> Tuple[str, Dict, List[Dict]]:
    """Async convenience function to generate podcast from JSON string"""
    config = PodcastConfig(
        output_dir=output_dir,
        audio_dir=os.path.join(output_dir, "audio"),
        image_dir=os.path.join(output_dir, "images")
    )
    generator = PodcastGenerator(config)
    return await generator.generate_podcast_from_json_async(json_string)


def generate_podcast_from_json(json_string: str, output_dir: str = "podcast_output") -> Tuple[str, Dict, List[Dict]]:
    """Convenience function to generate podcast from JSON string"""
    return asyncio.run(generate_podcast_from_json_async(json_string, output_dir))


# 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)
    print(json_string)
    
    try:
        audio_path, metadata, segment_results = generate_podcast_from_json(json_string)
        print(f"Generated podcast: {audio_path}")
        print(f"Metadata: {metadata}")
        print(f"Segment results: {segment_results}")
        return audio_path, metadata, segment_results
        
    except Exception as e:
        print(f"Error: {str(e)}")
        return None, None, None

