import os, openai, time, re
import pingcap.pe.flash.util as util
import botocore.exceptions


def raw_chatgpt(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "user", "content": line},
        ],
    )
    return response["choices"][0]["message"]["content"]


def raw_davinci(prompt, max_tokens=500):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.Completion.create(
        engine="text-davinci-003",
        prompt="\n\n" + prompt + "\n\n",
        temperature=0,
        max_tokens=max_tokens,
        top_p=1.0,
        frequency_penalty=0.0,
        presence_penalty=0.0,
    )
    return response.choices[0].text


def make_exam_question(subject, choice_count, correct_count):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are an instructor in distributed database and cloud computing technologies.",
            },
            {
                "role": "user",
                "content": "Create a multiple choice question regarding "
                + subject
                + " with "
                + str(choice_count)
                + " choices with "
                + str(correct_count)
                + " correct anwser.",
            },
        ],
    )
    results = response["choices"][0]["message"]["content"]
    return results


def refine_en_sentence_by_openai(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are an instructor in distributed database and cloud computing technologies.",
            },
            {
                "role": "assistant",
                "content": "`Translate` stands for translating a single line of text from one language to a single line of text in another language without expansion or explanation.",
            },
            {
                "role": "user",
                "content": "Refine the sentence between the sign =::= with better English: "
                + "=::="
                + line
                + "=::=",
            },
        ],
    )
    results = response["choices"][0]["message"]["content"].split("=::=")
    if len(results) == 1:
        return results[0]
    else:
        return results[1]


def refine_zh_sentence_by_openai(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are an instructor in distributed database and cloud computing technologies.",
            },
            {
                "role": "assistant",
                "content": "`Translate` stands for translating a single line of text from one language to a single line of text in another language without expansion or explanation.",
            },
            {
                "role": "user",
                "content": "Refine the sentence between the sign =::= with better Simplified Chinese, must not use Traditional Chinese: "
                + "=::="
                + line
                + "=::=",
            },
        ],
    )
    results = response["choices"][0]["message"]["content"].split("=::=")
    if len(results) == 1:
        return results[0]
    else:
        return results[1]


def refine_ja_sentence_by_openai(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are an instructor in distributed database and cloud computing technologies.",
            },
            {
                "role": "assistant",
                "content": "`Translate` stands for translating a single line of text from one language to a single line of text in another language without expansion or explanation.",
            },
            {
                "role": "user",
                "content": "Refine the sentence between the sign =::= with better 日本語, must not use Traditional Chinese: "
                + "=::="
                + line
                + "=::=",
            },
        ],
    )
    results = response["choices"][0]["message"]["content"].split("=::=")
    if len(results) == 1:
        return results[0]
    else:
        return results[1]


def translate_en2ja_by_openai(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    t = 1
    results = None
    while True:
        t += 1
        if t == 9:
            break
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[
                    {
                        "role": "system",
                        "content": "You are an instructor in distributed database and cloud computing technologies.",
                    },
                    {
                        "role": "system",
                        "content": "You are proficient in both English and Japanese. ",
                    },
                    {
                        "role": "assistant",
                        "content": "`Translate` stands for translating a single line of text from one language to a single line of text in another language without expansion or explanation.",
                    },
                    {
                        "role": "user",
                        "content": "Translate this after the marker `=::=` into Japanese =::= "
                        + line,
                    },
                ],
            )
            results = response["choices"][0]["message"]["content"].split("=::=")
        except openai.error.RateLimitError as rle:
            print("Error:\n", rle)
            print("Backoff and try again ... please wait")
            time.sleep(2**t)
    if len(results) == 1:
        return re.sub("=::= *", "", results[0])
    else:
        return re.sub("=::= *", "", results[1])


def translate_en2ja_by_openai_test(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    t = 1
    results = None
    while True:
        t += 1
        if t == 9:
            break
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[
                    {
                        "role": "assistant",
                        "content": "あなたはリレーショナル データベース管理システムおよび分散データベースの専門家です。",
                    },
                    {
                        "role": "assistant",
                        "content": "あなたはプロの翻訳家であり、校正者でもあります。原文の意味を変えずに自然な日本語に翻訳し、さらに日本語として自然な表現に校正してください。",
                    },
                    {
                        "role": "user",
                        "content": "これを `=::=` マーカーの後に日本語に翻訳 =::= " + line,
                    },
                ],
            )
            results = response["choices"][0]["message"]["content"].split("=::=")
        except openai.error.RateLimitError as rle:
            print("Error:\n", rle)
            print("Backoff and try again ... please wait")
            time.sleep(2**t)
    if len(results) == 1:
        return re.sub("=::= *", "", results[0])
    else:
        return re.sub("=::= *", "", results[1])


def convert_ja_to_desumasu(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    t = 1
    results = None
    final_result = None
    while True:
        t += 1
        if t == 9:
            break
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[
                    {
                        "role": "user",
                        "content": "Convert the sentence after the sign `=::=` to Japanese ですます体 in single line. Must not use 正しい敬語, use ですます体 is enough, and you must return the single line result only =::= "
                        + line,
                    },
                ],
            )
            results = response["choices"][0]["message"]["content"].split("=::=")
        except openai.error.RateLimitError as rle:
            print("Error:\n", rle)
            print("Backoff and try again ... please wait")
            time.sleep(2**t)
    if len(results) == 1:
        final_result = results[0]
    else:
        final_result = results[1]
    final_result = re.sub("^.*-> ", "", final_result)
    final_result = re.sub("^.*=> ", "", final_result)
    final_result = re.sub(".*「", "", final_result)
    final_result = re.sub("」.*", "", final_result)
    final_result = re.sub("』.*", "", final_result)
    final_result = re.sub(".*『", "", final_result)
    return final_result


def convert_ja_to_taigendome(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    t = 1
    results = None
    final_result = None
    while True:
        t += 1
        if t == 9:
            break
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[
                    {
                        "role": "user",
                        "content": "Convert the sentence after the sign `=::=` to Japanese 体言止め in single line. Must not use 正しい敬語, use 体言止め is enough, and you must return the single line result only =::= "
                        + line,
                    },
                ],
            )
            results = response["choices"][0]["message"]["content"].split("=::=")
        except openai.error.RateLimitError as rle:
            print("Error:\n", rle)
            print("Backoff and try again ... please wait")
            time.sleep(2**t)
    if len(results) == 1:
        final_result = results[0]
    else:
        final_result = results[1]
    final_result = re.sub("^.*-> ", "", final_result)
    final_result = re.sub("^.*=> ", "", final_result)
    final_result = re.sub(".*「", "", final_result)
    final_result = re.sub("」.*", "", final_result)
    final_result = re.sub("』.*", "", final_result)
    final_result = re.sub(".*『", "", final_result)
    return final_result


def convert_ja_to_past_tense(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    t = 1
    results = None
    final_result = None
    while True:
        t += 1
        if t == 9:
            break
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[
                    {
                        "role": "user",
                        "content": "Convert this Japanese sentence into past tense: "
                        + line,
                    },
                ],
            )
            results = response["choices"][0]["message"]["content"].split("=::=")
        except openai.error.RateLimitError as rle:
            print("Error:\n", rle)
            print("Backoff and try again ... please wait")
            time.sleep(2**t)
    if len(results) == 1:
        final_result = results[0]
    else:
        final_result = results[1]
    final_result = re.sub("^.*-> ", "", final_result)
    final_result = re.sub("^.*=> ", "", final_result)
    final_result = re.sub(".*「", "", final_result)
    final_result = re.sub("」.*", "", final_result)
    final_result = re.sub("』.*", "", final_result)
    final_result = re.sub(".*『", "", final_result)
    return final_result


def translate_en2ja_by_openai_for_srt(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    t = 1
    results = None
    while True:
        t += 1
        if t == 9:
            break
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=[
                    {
                        "role": "system",
                        "content": "You are an instructor in distributed database and cloud computing technologies.",
                    },
                    {
                        "role": "system",
                        "content": "You are proficient in both English and Japanese. ",
                    },
                    {
                        "role": "assistant",
                        "content": "`Translate` stands for translating a single line of text from one language to a single line of text in another language without expansion or explanation.",
                    },
                    {
                        "role": "user",
                        "content": "Translate this to Japanese and output in single line, you must think the sign |+| as invisible but you must put them back in the translated result if they exit in the original Engish sentence: "
                        + line,
                    },
                ],
            )
            results = response["choices"][0]["message"]["content"].split("=::=")
            if results != None:
                break
        except openai.error.RateLimitError as rle:
            print("Error:\n", rle)
            print("Backoff and try again ... please wait")
            time.sleep(2**t)
    if len(results) == 1:
        return re.sub("=::= *", "", results[0])
    else:
        return re.sub("=::= *", "", results[1])


def tranlsate_en2ja_by_aws_for_srt(line):
    session = util.get_boto_session()
    translate = session.client("translate")
    t = 1
    result = None
    while True:
        t += 1
        if t == 9:
            break
        try:
            response = translate.translate_text(
                Text=line,
                TerminologyNames=[
                    "terms-pe-en-ja",
                ],
                SourceLanguageCode="en",
                TargetLanguageCode="ja",
                Settings={"Formality": "INFORMAL"},
            )
            result = response["TranslatedText"]
            if result != None:
                break
            else:
                time.sleep(1)
                continue
        except botocore.exceptions.ClientError as error:
            print("Error:\n", error)
            print("Backoff and try again ... please wait")
            time.sleep(2**t)
            continue
    return result


def translate_en2zh_by_openai(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are an instructor in distributed database and cloud computing technologies.",
            },
            {
                "role": "assistant",
                "content": "`Translate` stands for translating a single line of text from one language to a single line of text in another language without expansion or explanation.",
            },
            {
                "role": "user",
                "content": "Translate this after the marker `=::=` into Simplied Chinese, and must not use Traditional Chinese =::= "
                + line,
            },
        ],
    )
    results = response["choices"][0]["message"]["content"].split("=::=")
    if len(results) == 1:
        return results[0]
    else:
        return results[1]


def translate_en2ja_by_aws(line):
    session = util.get_boto_session()
    translate = session.client("translate")
    t = 1
    result = None
    while True:
        t += 1
        if t == 9:
            break
        try:
            response = translate.translate_text(
                Text=line,
                TerminologyNames=[
                    "terms-pe-en-ja",
                ],
                SourceLanguageCode="en",
                TargetLanguageCode="ja",
                Settings={"Formality": "FORMAL"},
            )
            result = response["TranslatedText"]
            if result != None:
                break
            else:
                time.sleep(1)
                continue
        except botocore.exceptions.ClientError as error:
            print("Error:\n", error)
            print("Backoff and try again ... please wait")
            time.sleep(2**t)
            continue
    return result


def translate_zh2en_by_openai(line):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are an instructor in distributed database and cloud computing technologies.",
            },
            {
                "role": "assistant",
                "content": "`Translate` stands for translating a single line of text from one language to a single line of text in another language without expansion or explanation.",
            },
            {
                "role": "user",
                "content": "Translate this after the marker `=::=` into English =::= "
                + line,
            },
        ],
    )
    results = response["choices"][0]["message"]["content"].split("=::=")
    result = "N/A"
    if len(results) == 1:
        result = re.sub("=::= *", "", results[0])
    else:
        result = re.sub("=::= *", "", results[1])
    if '"' not in line:
        result = result.replace('"', "")
    return result


def translate_jakz2hirakana(line, max_tokens=10):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.Completion.create(
        engine="text-davinci-003",
        prompt="Translate this to 日本語 Hiragana:\n\n" + line + "\n\n",
        temperature=0,
        max_tokens=max_tokens,
        top_p=1.0,
        frequency_penalty=0.0,
        presence_penalty=0.0,
    )
    return response.choices[0].text


def wip(p: str) -> str:
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.Image.create(prompt=p, n=1, size="512x512")
    image_url = response["data"][0]["url"]
    return image_url