import openai
from pytube import YouTube
import argparse
import os
from tqdm import tqdm
from SRT import SRT_script
import stable_whisper
import whisper
from srt2ass import srt2ass

import subprocess

import time

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--link", help="youtube video link here", default=None, type=str, required=False)
    parser.add_argument("--video_file", help="local video path here", default=None, type=str, required=False)
    parser.add_argument("--audio_file", help="local audio path here", default=None, type=str, required=False)
    parser.add_argument("--srt_file", help="srt file input path here", default=None, type=str, required=False)  # New argument
    parser.add_argument("--download", help="download path", default='./downloads', type=str, required=False)
    parser.add_argument("--output_dir", help="translate result path", default='./results', type=str, required=False)
    parser.add_argument("--video_name", help="video name, if use video link as input, the name will auto-filled by youtube video name", default='placeholder', type=str, required=False)
    parser.add_argument("--model_name", help="model name only support gpt-4 and gpt-3.5-turbo", type=str, required=False, default="gpt-4") # default change to gpt-4
    parser.add_argument("-only_srt", help="set script output to only .srt file", action='store_true')
    parser.add_argument("-v", help="auto encode script with video", action='store_true')
    args = parser.parse_args()
    
    return args

def get_sources(args, download_path, result_path, video_name):
    # get source audio
    if args.link is not None and args.video_file is None:
        # Download audio from YouTube
        video_link = args.link
        video = None
        audio = None
        try:
            yt = YouTube(video_link)
            video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first()
            if video:
                video.download(f'{download_path}/video')
                print('Video download completed!')
            else:
                print("Error: Video stream not found")
            audio = yt.streams.filter(only_audio=True, file_extension='mp4').first()
            if audio:
                audio.download(f'{download_path}/audio')
                print('Audio download completed!')
            else:
                print("Error: Audio stream not found")
        except Exception as e:
            print("Connection Error")
            print(e) 
            exit()
        
        video_path = f'{download_path}/video/{video.default_filename}'
        audio_path = '{}/audio/{}'.format(download_path, audio.default_filename)
        audio_file = open(audio_path, "rb")
        if video_name == 'placeholder':
            video_name = audio.default_filename.split('.')[0]
    elif args.video_file is not None:
        # Read from local
        video_path = args.video_file

        if args.audio_file is not None:
            audio_file= open(args.audio_file, "rb")
            audio_path = args.audio_file
        else:
            output_audio_path = f'{download_path}/audio/{video_name}.mp3'
            subprocess.run(['ffmpeg', '-i', video_path, '-f', 'mp3', '-ab', '192000', '-vn', output_audio_path])
            audio_file = open(output_audio_path, "rb")
            audio_path = output_audio_path

    if not os.path.exists(f'{result_path}/{video_name}'):
        os.mkdir(f'{result_path}/{video_name}')

    if args.audio_file is not None:
        audio_file= open(args.audio_file, "rb")
        audio_path = args.audio_file
        pass

    return audio_path, audio_file, video_path, video_name

def get_srt_class(srt_file_en, result_path, video_name, audio_path, audio_file = None, whisper_model = 'base', method = "stable"):
    # Instead of using the script_en variable directly, we'll use script_input
    if srt_file_en is not None: 
        srt = SRT_script.parse_from_srt_file(srt_file_en)
    else:
        # using whisper to perform speech-to-text and save it in <video name>_en.txt under RESULT PATH.
        srt_file_en = "{}/{}/{}_en.srt".format(result_path, video_name, video_name)
        if not os.path.exists(srt_file_en):

            # use OpenAI API for transcribe
            if method == "api":
                transcript = openai.Audio.transcribe("whisper-1", audio_file) 

            # use local whisper model 
            elif method == "basic":
                model = whisper.load_model(whisper_model) # using base model in local machine (may use large model on our server)
                transcript = model.transcribe(audio_path)

            # use stable-whisper
            elif method == "stable":
                model = stable_whisper.load_model(whisper_model)
                transcript = model.transcribe(audio_path, regroup = False, initial_prompt="Hello, welcome to my lecture. Are you good my friend?")
                (
                    transcript
                    .split_by_punctuation(['.', '。', '?'])
                    .merge_by_gap(.15, max_words=3)
                    .merge_by_punctuation([' '])
                    .split_by_punctuation(['.', '。', '?'])
                )
                transcript = transcript.to_dict()
            else:
                raise ValueError("invalid speech to text method")

            srt = SRT_script(transcript['segments']) # read segments to SRT class

        else:
            srt = SRT_script.parse_from_srt_file(srt_file_en)
    return srt_file_en, srt

# Split the video script by sentences and create chunks within the token limit
def script_split(script_in, chunk_size = 1000):
    script_split = script_in.split('\n\n')
    script_arr = []
    range_arr = []
    start = 1
    end = 0
    script = ""
    for sentence in script_split:
        if len(script) + len(sentence) + 1 <= chunk_size:
            script += sentence + '\n\n'
            end+=1
        else:
            range_arr.append((start, end))
            start = end+1
            end += 1
            script_arr.append(script.strip())
            script = sentence + '\n\n'
    if script.strip():
        script_arr.append(script.strip())
        range_arr.append((start, len(script_split)-1))

    assert len(script_arr) == len(range_arr)
    return script_arr, range_arr

# check whether previous translation is done
# zh_file = "{}/{}/{}_zh.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME)
# segidx = 1
# if os.path.exists(zh_file):
#     temp_file = "{}/{}/temp.srt".format(RESULT_PATH, VIDEO_NAME)
#     if os.path.exists(temp_file):
#         os.remove(temp_file)
#     with open(zh_file, "r") as f0:
#         for count, _ in enumerate(f0):
#             pass
#         count += 1
#         segidx = int(count/4)+1
#     en_file = "{}/{}/{}_en.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME)
#     if args.srt_file is not None:
#         en_file = args.srt_file
#     with open(en_file, "r") as f1, open(temp_file, "a") as f2:
#         x = f1.readlines()
#         #print(len(x))
#         if count >= len(x):
#             print('Work already done! Please delete {}_zh.srt files in result directory first in order to rework'.format(VIDEO_NAME))
#             exit()
#         for i, line in enumerate(x):
#             if i >= count:
#                 f2.write(line)
                
#     srt = SRT_script.parse_from_srt_file(temp_file)
#     print('temp_contents')
#     print(srt.get_source_only())  


def get_response(model_name, sentence):
    if model_name == "gpt-3.5-turbo" or model_name == "gpt-4":
        response = openai.ChatCompletion.create(
            model=model_name,
            messages = [
                {"role": "system", "content": "You are a helpful assistant that translates English to Chinese and have decent background in starcraft2."},
                {"role": "system", "content": "Your translation has to keep the orginal format and be as accurate as possible."},
                {"role": "system", "content": "Your translation needs to be consistent with the number of sentences in the original."},
                {"role": "system", "content": "There is no need for you to add any comments or notes."},
                {"role": "user", "content": 'Translate the following English text to Chinese: "{}"'.format(sentence)}
            ],
            temperature=0.15
        )

        return response['choices'][0]['message']['content'].strip()
        
        
# Translate and save
def translate(srt, script_arr, range_arr, model_name, video_name, video_link):
    previous_length = 0
    for sentence, range in tqdm(zip(script_arr, range_arr)):
        # update the range based on previous length
        range = (range[0]+previous_length, range[1]+previous_length)

        # using chatgpt model
        print(f"now translating sentences {range}")
        flag = True
        while flag:
            flag = False
            try:
                translate = get_response(model_name, sentence)
            except Exception as e:
                print("An error has occurred during translation:",e)
                print("Retrying... the script will continue after 30 seconds.")
                time.sleep(30)
                flag = True
        srt.set_translation(translate, range, model_name, video_name, video_link)


def main():
    args = parse_args()

    # input check: input should be either video file or youtube video link.
    if args.link is None and args.video_file is None and args.srt_file is None and args.audio_file is None:
        print("need video source or srt file")
        exit()

    # set up
    openai.api_key = os.getenv("OPENAI_API_KEY")
    DOWNLOAD_PATH = args.download
    if not os.path.exists(DOWNLOAD_PATH):
        os.mkdir(DOWNLOAD_PATH)
        os.mkdir(f'{DOWNLOAD_PATH}/audio')
        os.mkdir(f'{DOWNLOAD_PATH}/video')

    RESULT_PATH = args.output_dir
    if not os.path.exists(RESULT_PATH):
        os.mkdir(RESULT_PATH)

    # set video name as the input file name if not specified
    if args.video_name == 'placeholder' :
        # set video name to upload file name
        if args.video_file is not None:
            VIDEO_NAME = args.video_file.split('/')[-1].split('.')[0]
        elif args.audio_file is not None:
            VIDEO_NAME = args.audio_file.split('/')[-1].split('.')[0]
        elif args.srt_file is not None:
            VIDEO_NAME = args.srt_file.split('/')[-1].split('.')[0].split("_")[0]
        else:
            VIDEO_NAME = args.video_name
    else:
        VIDEO_NAME = args.video_name

    audio_path, audio_file, video_path, VIDEO_NAME = get_sources(args, DOWNLOAD_PATH, RESULT_PATH, VIDEO_NAME)

    srt_file_en, srt = get_srt_class(args.srt_file, RESULT_PATH, VIDEO_NAME, audio_path, audio_file)

    # SRT class preprocess
    srt.form_whole_sentence()
    srt.spell_check_term()
    srt.correct_with_force_term()
    srt.write_srt_file_src(srt_file_en)
    script_input = srt.get_source_only()

    # write ass
    if not args.only_srt:
        assSub_en = srt2ass(srt_file_en, "default", "No", "Modest")
        print('ASS subtitle saved as: ' + assSub_en)

    script_arr, range_arr = script_split(script_input)

    translate(srt, script_arr, range_arr, args.model_name, VIDEO_NAME, args.link)

    # SRT post-processing
    srt.check_len_and_split()
    srt.remove_trans_punctuation()
    srt.write_srt_file_translate(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt")
    srt.write_srt_file_bilingual(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt")

    # write ass
    if not args.only_srt:
        assSub_zh = srt2ass(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", "default", "No", "Modest")
        print('ASS subtitle saved as: ' + assSub_zh)

    # encode to .mp4 video file
    if args.v:
        if args.only_srt:
            os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')
        else:
            os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.ass" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')


if __name__ == "__main__":
    main()