import os
import logging
import uuid
from typing import List, Tuple
import urllib.request
from PIL import Image
from gtts import gTTS
import moviepy.editor as mp
import gradio as gr
from hercai import Hercai

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('comic_generator.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class ComicVideoGenerator:
    def __init__(self):
        """Initialize comic video generator with Hercai API"""
        self.api = Hercai()
        self.image_size = (1024, 1024)
        
    def generate_comic_image(self, prompt: str) -> str:
        """Generate comic-style image from text prompt"""
        try:
            enhanced_prompt = (
                f"{prompt}, comic book style, vibrant colors, "
                "clear speech bubbles, dramatic lighting, "
                "detailed backgrounds, professional illustration"
            )
            
            result = self.api.draw_image(
                model="simurg",
                prompt=enhanced_prompt,
                negative_prompt="blurry, low quality, dark"
            )
            logger.info(f"Generated image for prompt: {prompt[:30]}...")
            return result["url"]
            
        except Exception as e:
            logger.error(f"Image generation failed: {e}")
            raise

    def process_image(self, url: str, save_path: str) -> str:
        """Download and process image to correct size"""
        try:
            urllib.request.urlretrieve(url, save_path)
            
            with Image.open(save_path) as img:
                img.thumbnail(self.image_size)
                new_img = Image.new('RGB', self.image_size, 'white')
                offset = tuple(map(lambda x, y: (x - y) // 2, 
                                 self.image_size, img.size))
                new_img.paste(img, offset)
                new_img.save(save_path, quality=95)
                
            logger.info(f"Processed and saved image: {save_path}")
            return save_path
            
        except Exception as e:
            logger.error(f"Image processing failed: {e}")
            raise

    def create_audio(self, text: str, save_path: str) -> str:
        """Generate audio narration from text"""
        try:
            tts = gTTS(text=text, lang='en')
            tts.save(save_path)
            logger.info(f"Created audio: {save_path}")
            return save_path
        except Exception as e:
            logger.error(f"Audio generation failed: {e}")
            raise

    def process_scene(self, prompt: str, scene_id: str) -> Tuple[str, str]:
        """Process single scene (image + audio)"""
        image_path = f"scene_{scene_id}.png"
        audio_path = f"audio_{scene_id}.mp3"
        
        image_url = self.generate_comic_image(prompt)
        image_file = self.process_image(image_url, image_path)
        audio_file = self.create_audio(prompt, audio_path)
        
        return image_file, audio_file

    def create_video(self, images: List[str], audios: List[str], 
                    output_path: str) -> str:
        """Create final video from images and audio"""
        try:
            clips = []
            for img, audio in zip(images, audios):
                audio_clip = mp.AudioFileClip(audio)
                video_clip = (mp.ImageClip(img)
                            .set_duration(audio_clip.duration)
                            .set_audio(audio_clip))
                clips.append(video_clip)

            final_clip = mp.concatenate_videoclips(clips)
            final_clip.write_videofile(
                output_path,
                fps=24,
                codec='libx264',
                audio_codec='aac'
            )
            logger.info(f"Created video: {output_path}")
            return output_path
            
        except Exception as e:
            logger.error(f"Video creation failed: {e}")
            raise

    def generate(self, text: str) -> str:
        """Main video generation pipeline"""
        try:
            scenes = [s.strip() for s in text.split(",,") if s.strip()]
            output_path = f"comic_{uuid.uuid4().hex[:8]}.mp4"
            
            images, audios = [], []
            for i, scene in enumerate(scenes):
                img, audio = self.process_scene(scene, f"{i}")
                images.append(img)
                audios.append(audio)
                
            return self.create_video(images, audios, output_path)
            
        except Exception as e:
            logger.error(f"Generation pipeline failed: {e}")
            raise

def create_interface():
    """Create Gradio interface"""
    generator = ComicVideoGenerator()
    
    examples = [
        "A magical forest at sunset.,, A brave knight finds a glowing crystal.,, The crystal transforms into a dragon.",
        "A busy city street.,, A mysterious package appears.,, The package opens to reveal a portal."
    ]
    
    with gr.Blocks(theme='default') as demo:
        gr.Markdown("# Comic Video Generator")
        
        with gr.Row():
            text_input = gr.Textbox(
                label="Story Text",
                placeholder="Enter story scenes separated by ',,'",
                lines=3
            )
        
        with gr.Row():
            generate_btn = gr.Button("Generate Comic")
            video_output = gr.Video(label="Generated Comic")
            
        gr.Examples(examples, text_input)
        
        generate_btn.click(
            fn=generator.generate,
            inputs=text_input,
            outputs=video_output
        )
        
    return demo

if __name__ == "__main__":
    interface = create_interface()
    interface.launch(debug=True)