import streamlit as st
from TTS.api import TTS
import torch
import subprocess
import locale
import assemblyai as aai
import requests

# Load TTS model
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=True).to("cuda")


# Function to translate video
class Translation:
    def __init__(self, video_path, original_language, target_language):
        self.video_path = video_path
        self.original_language = original_language
        self.target_language = target_language

    def org_language_parameters(self, original_language):
        if original_language == 'English':
            self.lan_code = 'en'
        elif original_language == 'German':
            self.lan_code = 'de'
        elif original_language == 'French':
            self.lan_code = 'fr'
        elif original_language == 'Spanish':
            self.lan_code = 'es'
        else:
            self.lan_code = ''

    def target_language_parameters(self, target_language):
        if target_language == 'English':
            self.tran_code = 'en'
        elif target_language == 'German':
            self.tran_code = 'de'
        elif target_language == 'French':
            self.tran_code = 'fr'
        elif target_language == 'Spanish':
            self.tran_code = 'es'
        else:
            self.tran_code = ''

    def extract_audio(self):
        audio_path = "output_audio.wav"
        subprocess.run(['ffmpeg', '-i', self.video_path, '-vn', '-acodec', 'pcm_s16le', '-ar', '44100', '-ac', '2', audio_path])
        st.success("Audio extracted successfully!")
        return audio_path

    def transcribe_audio(self, audio_path):
        aai.settings.api_key = "c29eb650444a4ae4bea01d82dd861cbb"
        config = aai.TranscriptionConfig(language_code=self.lan_code)
        transcriber = aai.Transcriber(config=config)
        transcript = transcriber.transcribe(audio_path)
        transcript_text = transcript.text
        return transcript_text

    def translate_text(self, transcript_text):
        base_url = "https://api.cognitive.microsofttranslator.com"
        endpoint = "/translate"
        headers = {
            "Ocp-Apim-Subscription-Key": "cd226bb1f3644276bea01d82dd861cbb",
            "Content-Type": "application/json",
            "Ocp-Apim-Subscription-Region": "southeastasia"
        }
        params = {
            "api-version": "3.0",
            "from": self.lan_code,
            "to": self.tran_code
        }
        body = [{"text": transcript_text}]
        response = requests.post(base_url + endpoint, headers=headers, params=params, json=body)
        response.raise_for_status()
        translation = response.json()[0]["translations"][0]["text"]
        return translation

    def generate_audio(self, translated_text):
        tts.tts_to_file(translated_text,
                        speaker_wav='output_audio.wav',
                        file_path="output_synth.wav",
                        language=self.tran_code
                        )
        return "output_synth.wav"

    def translate_video(self):
        audio_path = self.extract_audio()
        self.org_language_parameters(self.original_language)
        self.target_language_parameters(self.target_language)
        transcript_text = self.transcribe_audio(audio_path)
        translated_text = self.translate_text(transcript_text)
        translated_audio_path = self.generate_audio(translated_text)

        # Convert translated audio to video
        subprocess.run(['ffmpeg', '-i', self.video_path, '-i', translated_audio_path, '-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', 'output_video.mp4'])


st.title("Translate Your Video")
st.write("Upload your video and select the original and target languages.")

# Upload video
video_file = st.file_uploader("Upload Video", type=["mp4"])

if video_file is not None:
    # Get original and target languages
    original_language = st.selectbox("Select Original Language", ['English', 'German', 'French', 'Spanish'])
    target_language = st.selectbox("Select Target Language", ['English', 'German', 'French', 'Spanish'])

    translation = Translation(video_path=video_file.name,
                              original_language=original_language,
                              target_language=target_language)

    if st.button("Translate"):
        translation.translate_video()
        st.success("Video translation complete! You can download the translated video.")