import openai
from pytube import YouTube
import argparse
import os
from pathlib import Path
from tqdm import tqdm
from src.srt_util.srt import SrtScript
from src.Pigeon import Pigeon
import stable_whisper
import whisper
from src.srt_util import srt2ass
import logging
from datetime import datetime
import torch

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("--log_dir", help="log path", default='./logs', type=str, required=False)
    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
    audio_path = None
    audio_file = None
    video_path = None

    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,use_oauth=True, allow_oauth_cache=True)
            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='large',
                  method="stable"):
    # Instead of using the script_en variable directly, we'll use script_input
    if srt_file_en is not None:
        srt = SrtScript.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):

            devices = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            # 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,
                                           device=devices)  # using base model in local machine (may use large model on our server)
                transcript = model.transcribe(audio_path)
                srt = SRT_script(transcript['segments']) # read segments to SRT class

            # use stable-whisper
            elif method == "stable":

                # use cuda if available
                model = stable_whisper.load_model(whisper_model, device=devices)
                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()
                srt = SRT_script(transcript['segments']) # read segments to SRT class
            else:
                raise ValueError("invalid speech to text method")

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

        else:
            srt = SrtScript.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


def check_translation(sentence, translation):
    """
    check merge sentence issue from openai translation
    """
    sentence_count = sentence.count('\n\n') + 1
    translation_count = translation.count('\n\n') + 1

    if sentence_count != translation_count:
        # print("sentence length: ", len(sentence), sentence_count)
        # print("translation length: ",  len(translation), translation_count)
        return False
    else:
        return True


def get_response(model_name, sentence):
    """
    Generates a translated response for a given sentence using a specified OpenAI model.

    Args:
    model_name (str): The name of the OpenAI model to be used for translation, either "gpt-3.5-turbo" or "gpt-4".
    sentence (str): The English sentence related to StarCraft 2 videos that needs to be translated into Chinese.

    Returns:
    str: The translated Chinese sentence, maintaining the original format, meaning, and number of lines.
    """

    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)}

                {"role": "system",
                 "content": "你是一个翻译助理，你的任务是翻译星际争霸视频，你会被提供一个按行分割的英文段落，你需要在保证句意和行数的情况下输出翻译后的文本。"},
                {"role": "user", "content": 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, attempts_count=5):
    """
    Translates the given script array into another language using the chatgpt and writes to the SRT file.

    This function takes a script array, a range array, a model name, a video name, and a video link as input. It iterates
    through sentences and range in the script and range arrays. If the translation check fails for five times, the function
    will attempt to resolve merge sentence issues and split the sentence into smaller tokens for a better translation. 
    
    Args:
    srt (Subtitle): An instance of the Subtitle class representing the SRT file.
    script_arr (list): A list of strings representing the original script sentences to be translated.
    range_arr (list): A list of tuples representing the start and end positions of sentences in the script.
    model_name (str): The name of the translation model to be used.
    video_name (str): The name of the video.
    video_link (str): The link to the video.
    attempts_count (int): Number of attemps of failures for unmatched sentences.
    """

    logging.info("start translating...")
    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}")
        logging.info(f"now translating sentences {range}, time: {datetime.now()}")
        flag = True
        while flag:
            flag = False
            try:
                translate = get_response(model_name, sentence)
                # detect merge sentence issue and try to solve for five times:
                while not check_translation(sentence, translate) and attempts_count > 0:
                    translate = get_response(model_name, sentence)
                    attempts_count -= 1

                # if failure still happen, split into smaller tokens
                if attempts_count == 0:
                    single_sentences = sentence.split("\n\n")
                    logging.info("merge sentence issue found for range", range)
                    translate = ""
                    for i, single_sentence in enumerate(single_sentences):
                        if i == len(single_sentences) - 1:
                            translate += get_response(model_name, single_sentence)
                        else:
                            translate += get_response(model_name, single_sentence) + "\n\n"
                            # print(single_sentence, translate.split("\n\n")[-2])
                    logging.info("solved by individually translation!")

            except Exception as e:
                logging.debug("An error has occurred during translation:", 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_old():
    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:
        raise TypeError("need video source or srt file")

    # set up
    start_time = time.time()
    openai.api_key = os.getenv("OPENAI_API_KEY")
    DOWNLOAD_PATH = Path(args.download)
    if not DOWNLOAD_PATH.exists():
        DOWNLOAD_PATH.mkdir(parents=False, exist_ok=False)
        DOWNLOAD_PATH.joinpath('audio').mkdir(parents=False, exist_ok=False)
        DOWNLOAD_PATH.joinpath('video').mkdir(parents=False, exist_ok=False)

    RESULT_PATH = Path(args.output_dir)
    if not RESULT_PATH.exists():
        RESULT_PATH.mkdir(parents=False, exist_ok=False)

    # 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)

    if not os.path.exists(args.log_dir):
        os.makedirs(args.log_dir)
    logging.basicConfig(level=logging.INFO, handlers=[
        logging.FileHandler("{}/{}_{}.log".format(args.log_dir, VIDEO_NAME, datetime.now().strftime("%m%d%Y_%H%M%S")),
                            'w', encoding='utf-8')])
    logging.info("---------------------Video Info---------------------")
    logging.info("Video name: {}, translation model: {}, video link: {}".format(VIDEO_NAME, args.model_name, args.link))

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

    # SRT class preprocess
    logging.info("---------------------Start Preprocessing SRT class---------------------")
    srt.write_srt_file_src(srt_file_en)
    srt.form_whole_sentence()
    # srt.spell_check_term()
    # srt.correct_with_force_term()
    processed_srt_file_en = srt_file_en.split('.srt')[0] + '_processed.srt'
    srt.write_srt_file_src(processed_srt_file_en)
    script_input = srt.get_source_only()

    # write ass
    if not args.only_srt:
        logging.info("write English .srt file to .ass")
        assSub_en = srt2ass(processed_srt_file_en, "default", "No", "Modest")
        logging.info('ASS subtitle saved as: ' + assSub_en)

    script_arr, range_arr = script_split(script_input)
    logging.info("---------------------Start Translation--------------------")
    translate(srt, script_arr, range_arr, args.model_name, VIDEO_NAME, args.link)

    # SRT post-processing
    logging.info("---------------------Start Post-processing SRT class---------------------")
    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:
        logging.info("write Chinese .srt file to .ass")
        assSub_zh = srt2ass(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", "default", "No", "Modest")
        logging.info('ASS subtitle saved as: ' + assSub_zh)

    # encode to .mp4 video file
    if args.v:
        logging.info("encoding video file")
        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')

    end_time = time.time()
    logging.info(
        "Pipeline finished, time duration:{}".format(time.strftime("%H:%M:%S", time.gmtime(end_time - start_time))))


def main():
    pigeon = Pigeon()
    pigeon.run()


if __name__ == "__main__":
    main()