import streamlit as st
from moviepy.editor import VideoFileClip, AudioFileClip, TextClip, CompositeVideoClip, concatenate_audioclips
import whisper
from translate import Translator
from gtts import gTTS
import tempfile
import os
import numpy as np
from datetime import timedelta
import shutil
from pathlib import Path

# Set page configuration
st.set_page_config(
    page_title="Tamil Movie Dubber",
    page_icon="🎬",
    layout="wide"
)

# Custom CSS
st.markdown("""
    <style>
    .stButton>button {
        width: 100%;
        border-radius: 5px;
        height: 3em;
        background-color: #FF4B4B;
        color: white;
    }
    .stProgress .st-bo {
        background-color: #FF4B4B;
    }
    </style>
    """, unsafe_allow_html=True)

# Tamil voice configurations
TAMIL_VOICES = {
    'Female 1': {'name': 'ta-IN-PallaviNeural', 'style': 'normal'},
    'Female 2': {'name': 'ta-IN-PallaviNeural', 'style': 'formal'},
    'Male 1': {'name': 'ta-IN-ValluvarNeural', 'style': 'normal'},
    'Male 2': {'name': 'ta-IN-ValluvarNeural', 'style': 'formal'}
}

class TamilTextProcessor:
    @staticmethod
    def normalize_tamil_text(text):
        """Normalize Tamil text for better pronunciation"""
        tamil_numerals = {'௦': '0', '௧': '1', '௨': '2', '௩': '3', '௪': '4',
                         '௫': '5', '௬': '6', '௭': '7', '௮': '8', '௯': '9'}
        for tamil_num, eng_num in tamil_numerals.items():
            text = text.replace(tamil_num, eng_num)
        return text

    @staticmethod
    def process_for_tts(text):
        """Process Tamil text for TTS"""
        text = ''.join(char for char in text if ord(char) < 65535)
        text = ' '.join(text.split())
        return text

@st.cache_resource
def load_whisper_model():
    """Load Whisper model with caching"""
    return whisper.load_model("base")

class VideoProcessor:
    def __init__(self):
        self.temp_dir = Path(tempfile.mkdtemp())
        self.whisper_model = load_whisper_model()

    def create_temp_path(self, suffix):
        """Create a temporary file path"""
        return str(self.temp_dir / f"temp_{os.urandom(4).hex()}{suffix}")

    def cleanup(self):
        """Clean up temporary directory"""
        try:
            shutil.rmtree(self.temp_dir)
        except Exception as e:
            st.warning(f"Cleanup warning: {e}")

    def transcribe_video(self, video_path):
        """Transcribe video audio using Whisper"""
        try:
            with VideoFileClip(video_path) as video:
                # Extract audio to temporary file
                audio_path = self.create_temp_path(".wav")
                video.audio.write_audiofile(audio_path, fps=16000, verbose=False, logger=None)

                # Check if audio file is not empty
                if os.path.getsize(audio_path) == 0:
                    raise ValueError("Extracted audio file is empty")

                # Transcribe using Whisper
                result = self.whisper_model.transcribe(audio_path)
                return result["segments"], video.duration

        except Exception as e:
            raise Exception(f"Transcription error: {str(e)}")

    def translate_segments(self, segments):
        """Translate segments to Tamil"""
        translator = Translator(to_lang='ta')
        translated_segments = []
        
        for segment in segments:
            try:
                translated_text = translator.translate(segment["text"])
                translated_text = TamilTextProcessor.normalize_tamil_text(translated_text)
                translated_text = TamilTextProcessor.process_for_tts(translated_text)
                
                translated_segments.append({
                    "text": translated_text,
                    "start": segment["start"],
                    "end": segment["end"],
                    "duration": segment["end"] - segment["start"]
                })
            except Exception as e:
                st.warning(f"Translation warning for segment: {str(e)}")
                # Keep original text if translation fails
                translated_segments.append({
                    "text": segment["text"],
                    "start": segment["start"],
                    "end": segment["end"],
                    "duration": segment["end"] - segment["start"]
                })
        
        return translated_segments

    def generate_tamil_audio(self, text):
        """Generate Tamil audio using gTTS"""
        try:
            audio_path = self.create_temp_path(".mp3")
            tts = gTTS(text=text, lang='ta', slow=False)
            tts.save(audio_path)
            return audio_path
        except Exception as e:
            raise Exception(f"Audio generation error: {str(e)}")

    def create_subtitle_clip(self, txt, fontsize, color, size):
        """Create a subtitle clip"""
        return TextClip(
            txt=txt,
            fontsize=fontsize,
            color=color,
            bg_color='rgba(0,0,0,0.5)',
            size=size,
            method='caption'
        )

def process_video(video_data, voice_type, generate_subtitles=True, subtitle_size=24, subtitle_color='white'):
    """Main video processing function"""
    processor = VideoProcessor()
    
    try:
        # Save uploaded video to temporary file
        input_path = processor.create_temp_path(".mp4")
        with open(input_path, "wb") as f:
            f.write(video_data)

        # Load video
        video = VideoFileClip(input_path)
        
        # Create progress tracking
        progress_text = st.empty()
        progress_bar = st.progress(0)
        
        # Step 1: Transcribe
        progress_text.text("Transcribing video...")
        segments, duration = processor.transcribe_video(input_path)
        progress_bar.progress(0.25)
        
        # Step 2: Translate
        progress_text.text("Translating to Tamil...")
        translated_segments = processor.translate_segments(segments)
        progress_bar.progress(0.50)
        
        # Step 3: Generate audio
        progress_text.text("Generating Tamil audio...")
        subtitle_clips = []
        audio_clips = []
        
        for i, segment in enumerate(translated_segments):
            # Generate audio
            audio_path = processor.generate_tamil_audio(segment["text"])
            audio_clip = AudioFileClip(audio_path)
            audio_clips.append(audio_clip.set_start(segment["start"]))
            
            # Create subtitle if enabled
            if generate_subtitles:
                subtitle_clip = processor.create_subtitle_clip(
                    segment["text"],
                    subtitle_size,
                    subtitle_color,
                    (video.w, None)
                )
                subtitle_clip = (subtitle_clip
                               .set_position(('center', 'bottom'))
                               .set_start(segment["start"])
                               .set_duration(segment["duration"]))
                subtitle_clips.append(subtitle_clip)
            
            progress_bar.progress(0.50 + (0.4 * (i + 1) / len(translated_segments)))
        
        # Step 4: Combine everything
        progress_text.text("Creating final video...")
        
        # Combine audio clips
        final_audio = concatenate_audioclips(audio_clips)
        
        # Create final video
        if generate_subtitles:
            final_video = CompositeVideoClip([video, *subtitle_clips])
        else:
            final_video = video
        
        # Set audio
        final_video = final_video.set_audio(final_audio)
        
        # Write final video
        output_path = processor.create_temp_path(".mp4")
        final_video.write_videofile(
            output_path,
            codec='libx264',
            audio_codec='aac',
            temp_audiofile=processor.create_temp_path(".m4a"),
            remove_temp=True,
            verbose=False,
            logger=None
        )
        
        progress_bar.progress(1.0)
        progress_text.text("Processing complete!")
        
        return output_path
        
    except Exception as e:
        raise Exception(f"Video processing error: {str(e)}")
        
    finally:
        # Cleanup
        processor.cleanup()

def main():
    st.title("Tamil Movie Dubbing System")
    st.markdown("""
    👋 Welcome! This tool helps you:
    - 🎥 Convert English videos to Tamil
    - 🗣️ Generate Tamil voiceovers
    - 📝 Add Tamil subtitles
    """)
    
    # File uploader
    video_file = st.file_uploader("Upload Video File", type=['mp4', 'mov', 'avi'])
    
    if not video_file:
        st.warning("Please upload a video to begin.")
        return
    
    # Settings
    col1, col2 = st.columns(2)
    
    with col1:
        voice_type = st.selectbox("Select Voice", list(TAMIL_VOICES.keys()))
    
    with col2:
        generate_subtitles = st.checkbox("Generate Subtitles", value=True)
        
    # Ensure subtitle size and color are always defined
    subtitle_size = 24  # Default subtitle size
    subtitle_color = 'white'  # Default subtitle color
    
    if generate_subtitles:
        subtitle_size = st.slider("Subtitle Size", 10, 40, 24)
        subtitle_color = st.color_picker("Subtitle Color", "#FFFFFF")

    if st.button("Process Video"):
        try:
            output_video_path = process_video(
                video_file.read(),
                voice_type,
                generate_subtitles=generate_subtitles,
                subtitle_size=subtitle_size,
                subtitle_color=subtitle_color
            )
            st.success("Video processing complete!")
            st.video(output_video_path)
        except Exception as e:
            st.error(f"Error: {e}")

if __name__ == "__main__":
    main()