import os
import asyncio
import re
import numpy as np
import whisper
import edge_tts
import gradio as gr
from moviepy import VideoFileClip, TextClip, CompositeVideoClip, AudioFileClip, vfx, concatenate_audioclips, AudioArrayClip, CompositeAudioClip

# TikTok-style Pop-up Animation
def apply_tiktok_animation(clip, duration=0.2):
    def zoom(t):
        if t < duration / 2:
            return 0.5 + 1.4 * (t / duration)
        elif t < duration:
            return 1.2 - 0.4 * ((t - duration/2) / (duration/2))
        return 1.0
    return clip.with_effects([vfx.Resize(zoom)])

async def generate_video(text, video_path, voice, voice_speed, font_size, font_color, stroke_width, stroke_color, progress=gr.Progress()):
    try:
        # 1. Generate Speech
        progress(0, desc="📢 Generating speech...")
        parts = re.split(r'\[PAUSE\s+(\d+(?:\.\d+)?)\]', text)
        audio_clips = []
        temp_files = []
        for i, part in enumerate(parts):
            if i % 2 == 0:
                if part.strip():
                    chunk_file = f"temp_chunk_{i}.mp3"
                    communicate = edge_tts.Communicate(part.strip(), voice, rate=voice_speed)
                    await communicate.save(chunk_file)
                    audio_clips.append(AudioFileClip(chunk_file))
                    temp_files.append(chunk_file)
            else:
                silence = AudioArrayClip(np.zeros((int(44100 * float(part)), 2)), fps=44100)
                audio_clips.append(silence)

        combined_audio = concatenate_audioclips(audio_clips)
        
        # 2. Process Video
        progress(0.2, desc="🎞️ Loading and trimming video...")
        video_clip = VideoFileClip(video_path)
        if video_clip.duration < combined_audio.duration:
            final_video_clip = video_clip.with_effects([vfx.Loop(duration=combined_audio.duration)])
        else:
            final_video_clip = video_clip.subclipped(0, combined_audio.duration)

        if final_video_clip.audio is not None:
            bg_audio = final_video_clip.audio.with_volume_scaled(0.5)
            final_video_clip = final_video_clip.with_audio(CompositeAudioClip([bg_audio, combined_audio]))
        else:
            final_video_clip = final_video_clip.with_audio(combined_audio)

        # 3. Subtitles
        progress(0.4, desc="📝 Transcribing for subtitles...")
        temp_audio = "temp_full.mp3"
        combined_audio.write_audiofile(temp_audio, logger=None)
        
        # We use 'base' for local, but you can change back to 'tiny' for Space if needed
        model = whisper.load_model("tiny") 
        result = model.transcribe(temp_audio, word_timestamps=True)
        
        progress(0.6, desc="✨ Creating animated captions...")
        subtitle_clips = []
        for segment in result['segments']:
            for word_info in segment.get('words', []):
                word_text = word_info['word'].strip().upper()
                
                # New "Neon Glow" Style
                # We set these to start at 0 internally so the container handles the real timing
                duration = word_info['end'] - word_info['start']
                if duration <= 0: duration = 0.1

                glow_clip = TextClip(
                    text=word_text,
                    font_size=font_size, 
                    color=font_color,
                    stroke_color=font_color,
                    stroke_width=stroke_width + 6,
                    method='caption', 
                    size=(int(final_video_clip.w * 0.9), None)
                ).with_duration(duration).with_position('center').with_opacity(0.4)
                
                txt_clip = TextClip(
                    text=word_text,
                    font_size=font_size, 
                    color=font_color,
                    stroke_color=stroke_color, 
                    stroke_width=stroke_width,
                    method='caption', 
                    size=(int(final_video_clip.w * 0.9), None)
                ).with_duration(duration).with_position('center')
                
                # Combine layers into one "word unit" starting at 0 internally
                word_unit = CompositeVideoClip([glow_clip, txt_clip], size=txt_clip.size, bg_color=None).with_duration(duration)
                
                # Now apply the animation and the REAL start time to the whole unit
                word_unit = apply_tiktok_animation(word_unit)
                word_unit = word_unit.with_start(word_info['start']).with_position(('center', int(final_video_clip.h * 0.65)))
                
                subtitle_clips.append(word_unit)

        # 4. Export
        progress(0.8, desc="🚀 Rendering final video (Exporting)...")
        output_path = "output_video.mp4"
        final_video = CompositeVideoClip([final_video_clip] + subtitle_clips)
        final_video.write_videofile(output_path, codec="libx264", audio_codec="aac", fps=24, logger=None)

        progress(1.0, desc="✅ Complete!")
        
        # Final cleanup of temp audio
        if os.path.exists(temp_audio): os.remove(temp_audio)
        for f in temp_files:
            if os.path.exists(f): os.remove(f)

        return output_path, "Done!"
    except Exception as e:
        import traceback
        return None, f"Error: {str(e)}\n{traceback.format_exc()}"

# UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎬 tts-captions-app PRO")
    gr.Markdown("Generating videos with TikTok animations and Glow subtitles.")
    
    with gr.Row():
        with gr.Column():
            txt = gr.Textbox(label="Text", value="WELCOME! [PAUSE 0.5] This app has a GLOW effect [PAUSE 0.3] and a progress bar!!", lines=4)
            vid = gr.File(label="Background Video (Drag & Drop)", type="filepath")
            
            with gr.Accordion("Settings", open=False):
                voice = gr.Dropdown(label="Voice", choices=["en-US-ChristopherNeural", "en-US-AriaNeural", "en-GB-RyanNeural"], value="en-US-ChristopherNeural")
                size = gr.Slider(label="Font Size", minimum=20, maximum=200, value=80)
                color = gr.ColorPicker(label="Text Color", value="#FFFFFF")
                s_color = gr.ColorPicker(label="Outline Color", value="#000000")
                s_width = gr.Slider(label="Outline Width", minimum=0, maximum=20, value=10)
            
            btn = gr.Button("🚀 Generate Video", variant="primary")
        
        with gr.Column():
            out_v = gr.Video(label="Result")
            out_t = gr.Textbox(label="Log / Error Console")

    # Linking the function with Gradio's Progress feature
    btn.click(
        fn=generate_video, 
        inputs=[txt, vid, voice, gr.State("+20%"), size, color, s_width, s_color], 
        outputs=[out_v, out_t]
    )

if __name__ == "__main__":
    demo.launch()