from datetime import timedelta
import os
import subprocess
import whisper
import tempfile
import argparse
import langchain
from langchain.chat_models import ChatOpenAI, ChatGooglePalm
from langchain.schema import HumanMessage, SystemMessage, AIMessage
from langchain.prompts import (
    ChatPromptTemplate,
    PromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.chains import LLMChain
from langchain.callbacks import get_openai_callback
from tqdm import tqdm

def get_translate_chain(from_lang, to_lang):
    template=f"You are a helpful assistant that translates {from_lang} to {to_lang}."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_template="Please translate \"{text}\""+f" from {from_lang} to {to_lang}. Give me the translated {to_lang} directly without saying anything else, do not use \"."
    human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])

    # get a chat completion from the formatted messages
    chat = ChatOpenAI()
    chain = LLMChain(llm=chat, prompt=chat_prompt, verbose=True)
    return chain


def gen_srt(video_path, model_name="medium", from_language="English", to_language="Chinese", embed=False, translate=True):
    with tempfile.TemporaryDirectory() as temp_dir:
        # 1. use ffmpeg to extract audio from video and save it to Temp folder
        # Path to the temporary audio file
        temp_audio_path = os.path.join(temp_dir, "extracted_audio.wav")
        
        # Use ffmpeg to extract audio from video
        print("Extracting audio from video...")
        command = f"ffmpeg -i {video_path} -vn -ar 44100 -ac 2 -b:a 192k {temp_audio_path}"
        
        # Execute the command
        subprocess.call(command, shell=True)

        model = whisper.load_model(model_name)
        transcribe = model.transcribe(audio=temp_audio_path, language=from_language)
        segments = transcribe['segments']
        
        # 2. Use whisper to transcribe audio and save segments to srt file
        if translate:
            with get_openai_callback() as cb:
                chain = get_translate_chain(from_language, to_language)
                for segment in tqdm(segments):
                    segment['text'] = chain(segment['text'])['text']
                print(cb)
        
        # 3. Generate the SRT file
        srtFilename = video_path.split(".")[0] + ".srt"
        # overwrite the file if it already exists
        if os.path.exists(srtFilename):
            os.remove(srtFilename)
        for segment in segments:
            startTime = str(0)+str(timedelta(seconds=int(segment['start'])))+',000'
            endTime = str(0)+str(timedelta(seconds=int(segment['end'])))+',000'
            text = segment['text']
            segmentId = segment['id']+1
            segment = f"{segmentId}\n{startTime} --> {endTime}\n{text[1:] if text[0] == ' ' else text}\n\n"

            with open(srtFilename, 'a', encoding='utf-8') as srtFile:
                srtFile.write(segment)
        
        # 4. Use FFMPEG to embed srt file into video
        if not embed:
            return
        output_filename = video_path.split(".")[0] + "_subtitled.mp4"
        if os.path.exists(output_filename):
            os.remove(output_filename)
        embed_command = f"ffmpeg -i {video_path} -vf subtitles={srtFilename} {output_filename}"
        subprocess.call(embed_command, shell=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process some arguments')

    # Add the arguments
    parser.add_argument('-i', type=str, required=True, dest='input_file',
                        help='Input file name')

    parser.add_argument('-m', type=str, default='medium', dest='model_name',
                        help='Model type, default is "medium"')

    parser.add_argument('-f', type=str, default='English', dest='from_lang',
                        help='Translate from language, default is "English"')

    parser.add_argument('-t', type=str, default='Chinese', dest='to_lang',
                        help='Translate to language, default is "Chinese"')

    parser.add_argument('--embed', dest='embed', action='store_true',
                        help='Whether to Embed subtitles, default is False')

    parser.add_argument('--translate', dest='translate', action='store_true',
                        help='Whether to Translate, default is False')
    args = parser.parse_args()

    gen_srt(args.input_file, model_name=args.model_name, embed=args.embed, translate=args.translate, from_language=args.from_lang, to_language=args.to_lang)

langchain.prompts.HumanMessagePromptTemplate.from_template langchain.chat_models.ChatOpenAI langchain.prompts.ChatPromptTemplate.from_messages langchain.callbacks.get_openai_callback langchain.chains.LLMChain langchain.prompts.SystemMessagePromptTemplate.from_template
[((696, 747), 'langchain.prompts.SystemMessagePromptTemplate.from_template', 'SystemMessagePromptTemplate.from_template', (['template'], {}), '(template)\n', (737, 747), False, 'from langchain.prompts import ChatPromptTemplate, PromptTemplate, SystemMessagePromptTemplate, AIMessagePromptTemplate, HumanMessagePromptTemplate\n'), ((946, 1002), 'langchain.prompts.HumanMessagePromptTemplate.from_template', 'HumanMessagePromptTemplate.from_template', (['human_template'], {}), '(human_template)\n', (986, 1002), False, 'from langchain.prompts import ChatPromptTemplate, PromptTemplate, SystemMessagePromptTemplate, AIMessagePromptTemplate, HumanMessagePromptTemplate\n'), ((1021, 1100), 'langchain.prompts.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (['[system_message_prompt, human_message_prompt]'], {}), '([system_message_prompt, human_message_prompt])\n', (1053, 1100), False, 'from langchain.prompts import ChatPromptTemplate, PromptTemplate, SystemMessagePromptTemplate, AIMessagePromptTemplate, HumanMessagePromptTemplate\n'), ((1169, 1181), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {}), '()\n', (1179, 1181), False, 'from langchain.chat_models import ChatOpenAI, ChatGooglePalm\n'), ((1194, 1246), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'chat', 'prompt': 'chat_prompt', 'verbose': '(True)'}), '(llm=chat, prompt=chat_prompt, verbose=True)\n', (1202, 1246), False, 'from langchain.chains import LLMChain\n'), ((3618, 3679), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Process some arguments"""'}), "(description='Process some arguments')\n", (3641, 3679), False, 'import argparse\n'), ((1398, 1427), 'tempfile.TemporaryDirectory', 'tempfile.TemporaryDirectory', ([], {}), '()\n', (1425, 1427), False, 'import tempfile\n'), ((1589, 1634), 'os.path.join', 'os.path.join', (['temp_dir', '"""extracted_audio.wav"""'], {}), "(temp_dir, 'extracted_audio.wav')\n", (1601, 1634), False, 'import os\n'), ((1880, 1916), 'subprocess.call', 'subprocess.call', (['command'], {'shell': '(True)'}), '(command, shell=True)\n', (1895, 1916), False, 'import subprocess\n'), ((1934, 1964), 'whisper.load_model', 'whisper.load_model', (['model_name'], {}), '(model_name)\n', (1952, 1964), False, 'import whisper\n'), ((2619, 2646), 'os.path.exists', 'os.path.exists', (['srtFilename'], {}), '(srtFilename)\n', (2633, 2646), False, 'import os\n'), ((3357, 3388), 'os.path.exists', 'os.path.exists', (['output_filename'], {}), '(output_filename)\n', (3371, 3388), False, 'import os\n'), ((3533, 3575), 'subprocess.call', 'subprocess.call', (['embed_command'], {'shell': '(True)'}), '(embed_command, shell=True)\n', (3548, 3575), False, 'import subprocess\n'), ((2660, 2682), 'os.remove', 'os.remove', (['srtFilename'], {}), '(srtFilename)\n', (2669, 2682), False, 'import os\n'), ((3402, 3428), 'os.remove', 'os.remove', (['output_filename'], {}), '(output_filename)\n', (3411, 3428), False, 'import os\n'), ((2215, 2236), 'langchain.callbacks.get_openai_callback', 'get_openai_callback', ([], {}), '()\n', (2234, 2236), False, 'from langchain.callbacks import get_openai_callback\n'), ((2347, 2361), 'tqdm.tqdm', 'tqdm', (['segments'], {}), '(segments)\n', (2351, 2361), False, 'from tqdm import tqdm\n')]