import argparse
import json
import math
import os
import re
from pathlib import Path
from time import sleep

import tiktoken
from lmwrapper.openai_wrapper import OpenAiModelNames, OpenAIPredictor, get_open_ai_lm
from lmwrapper.structs import LmPrompt

# by default should be defects4j. In some cases, it might be <PATH_TO_DEFECTS4J>/framework/bin/defects4j
DEFECTS4J_COMMAND = "defects4j"


def preprocess(data):
    # creates the Defects4J_projects directory and checks out the buggy and fixed versions of the project

    print("Checking out buggy and fixed versions of projects...")
    # check if the Defects4J_projects directory exists
    if not Path("Defects4J_projects").exists():
        # if it doesn't exist, create it
        Path("Defects4J_projects").mkdir()

    for project in data:
        print(project)
        # this project is not deprecated
        # check if the buggy version of the project exists
        if not Path("Defects4J_projects/" + project[0] + "_" + project[1]).exists():
            # if it doesn't exist, check it out
            os.system(
                DEFECTS4J_COMMAND
                + " checkout -p "
                + project[0]
                + " -v "
                + project[1]
                + "b -w Defects4J_projects/"
                + project[0]
                + "_"
                + project[1]
                + "_buggy",
            )

        # check if the fixed version of the project exists
        if not Path(
            "Defects4J_projects/" + project[0] + "_" + project[1] + "_fixed",
        ).exists():
            # if it doesn't exist, check it out
            os.system(
                DEFECTS4J_COMMAND
                + " checkout -p "
                + project[0]
                + " -v "
                + project[1]
                + "f -w Defects4J_projects/"
                + project[0]
                + "_"
                + project[1]
                + "_fixed",
            )

        # copy over .defects4j.config from the fixed version to the buggy version
        # this is done because of an edge case where the fixed version removes some tests that are present in the buggy version
        # this causes defects4j to fail when running the test suite even though the buggy version should be passing
        os.system(
            "cp Defects4J_projects/"
            + project[0]
            + "_"
            + project[1]
            + "_fixed/.defects4j.config Defects4J_projects/"
            + project[0]
            + "_"
            + project[1]
            + "_buggy/.defects4j.config",
        )


def test_passing_data(project):
    # save the current working directory
    cwd = os.getcwd()

    # navigate to the buggy version of the project
    os.chdir("Defects4J_projects/" + project[0] + "_" + project[1] + "_buggy")

    try:
        # this block indicates a successful code change
        # run the test suite and pipe the output to a file
        os.system(
            DEFECTS4J_COMMAND
            + " test > "
            + project[0]
            + "_"
            + project[1]
            + "_buggy_test_output.txt",
        )

        # open the file
        with open(
            project[0] + "_" + project[1] + "_buggy_test_output.txt",
            encoding="locale",
        ) as file:
            # extract the number of failing tests. The number of failing tests is after the phrase "Failing tests: "
            failing_tests = file.read().split("Failing tests: ")[1].split("\n")[0]
            print("Failing tests: " + failing_tests)

        # get the total number of tests
        # this is the length of the file "all_tests"
        with open("all_tests", encoding="locale") as file:
            total_tests = len(file.readlines())
            print("Total tests: " + str(total_tests))

        passing_tests = total_tests - int(failing_tests)

    except:
        # this block indicates a compilation error
        # we go into this block if the test suite fails to run

        # navigate into the fixed version of the project
        os.chdir("../" + project[0] + "_" + project[1] + "_fixed")
        os.system(
            DEFECTS4J_COMMAND
            + " test > "
            + project[0]
            + "_"
            + project[1]
            + "_fixed_test_output.txt",
        )

        passing_tests = 0

        with open("all_tests", encoding="locale") as file:
            total_tests = str(len(file.readlines()))
            print("Total tests: " + total_tests)

    # navigate back to the original directory
    os.chdir(cwd)

    return passing_tests, total_tests


def check_equal_sign(line):
    return (
        "=" in line
        and "==" not in line
        and "!=" not in line
        and "<=" not in line
        and ">=" not in line
        and "=>" not in line
    )


def main():
    parser = argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--model",
        default=None,
        type=str,
        required=True,
        help="codex/turbo",
    )

    args = parser.parse_args()

    if args.model == "codex":
        lm = get_open_ai_lm(
            OpenAiModelNames.code_davinci_002,
        )
        OpenAIPredictor.configure_global_ratelimit(20, per_seconds=60)
    if args.model == "turbo":
        lm = get_open_ai_lm(
            OpenAiModelNames.gpt_3_5_turbo_instruct,
        )

    import csv

    # Set a larger field size limit, e.g., 2 MB (you can adjust this value as needed)
    csv.field_size_limit(10 * 1024 * 1024)

    # Specify the file path
    file_path = "Defects4J_oneLiner_metadata.csv"

    # Initialize an empty list to store the data
    data = []

    # Open the CSV file and read its contents
    with open(file_path, newline="", encoding="locale") as csvfile:
        csvreader = csv.reader(csvfile)

        # Loop through each row in the CSV file
        for row in csvreader:
            data.append(row)

    preprocess(data)

    ff = open(args.model + "_result_log_prob.txt", "a", encoding="locale")
    ff.write(
        "buggy"
        + "\t\t"
        + "Fixed"
        + "\t\t"
        + "generated_text"
        + "\t\t"
        + "is_exact"
        + "\t\t"
        + "total_prob"
        + "\t\t"
        + "average_prob"
        + "\t\t"
        + "log_probs"
        + "\t\t"
        + "length"
        + "\t\t"
        + "total_tests"
        + "\t\t"
        + "passing_tests"
        + "\n",
    )
    ff.close()

    enc = tiktoken.encoding_for_model("code-davinci-002")

    error_count = 0
    i = 0
    correct = 0
    json_objects = []
    while i < len(data):
        try:
            print(i)
            print(data[i])
            is_exact = 0
            file_name = data[i][2]
            line_number = int(data[i][3])

            # create a Path of Defects4J_projects/data[i][0]_data[i][1]_buggy
            buggy_file = (
                Path("Defects4J_projects/" + data[i][0] + "_" + data[i][1] + "_buggy")
                / file_name
            )
            fixed_file = (
                Path("Defects4J_projects/" + data[i][0] + "_" + data[i][1] + "_fixed")
                / file_name
            )

            print("Buggy file: ", buggy_file)
            print("Fixed file: ", fixed_file)

            # navigate to the buggy file and print out the contents of line at line_number
            with open(buggy_file, encoding="locale") as file:
                buggy_file_contents = file.readlines()
                buggy_line = buggy_file_contents[line_number - 1]
                print("Buggy line: ", buggy_line)

                # create a prompt for the buggy line
                # #############################
                before_buggy_line = buggy_file_contents[:line_number]
                after_buggy_line = buggy_file_contents[line_number:]
                # take last element of list split it by "=" and add <buggy> tag to the last element
                # only if the line contains "=" but not "==", "!=", "<=", ">=" and not "=>"
                if check_equal_sign(before_buggy_line[-1]):
                    before_buggy_line[-1] = (
                        before_buggy_line[-1].strip().split("=")[0]
                        + "= "
                        + "<buggy>"
                        + before_buggy_line[-1].strip().split("=")[1]
                        + "</buggy>"
                    )
                else:
                    before_buggy_line[-1] = (
                        "<buggy>" + before_buggy_line[-1] + "</buggy>"
                    )
                # convert list to string
                before_buggy_line = "".join(before_buggy_line)
                before_buggy_line = enc.encode(before_buggy_line)
                before_buggy_line = before_buggy_line[
                    len(before_buggy_line) - 200 : len(before_buggy_line)
                ]
                before_buggy_line = enc.decode(before_buggy_line)
                # convert list to string
                after_buggy_line = "".join(after_buggy_line)
                after_buggy_line = enc.encode(after_buggy_line)
                after_buggy_line = after_buggy_line[0:200]
                after_buggy_line = enc.decode(after_buggy_line)
                # join the two strings
                prompt = before_buggy_line + after_buggy_line + "\n"
                json_objects.append(
                    {
                        "project": data[i][0],
                        "number": data[i][1],
                        "prompt": prompt,
                    },
                )
                prompt += (
                    "Question: There is a bug in the above code snippet tagged by"
                    " <buggy> and </buggy>. Please generate the correct version.\n"
                )
                before_fixed_line = before_buggy_line.split("<buggy>")[0] + "<fixed>"
                prompt += "Answer:\n" + before_fixed_line
                # #############################

            # navigate to the fixed file and print out the contents of line at line_number
            with open(fixed_file, encoding="locale") as file:
                fixed_file_contents = file.readlines()
                fixed_line = fixed_file_contents[line_number - 1]

            print("Buggy line: ", buggy_line)
            print("Fixed line: ", fixed_line)
            print("Prompt: ", prompt)

            pred = lm.predict(
                LmPrompt(
                    text=prompt,
                    max_tokens=200,
                    logprobs=5,
                    stop=["\n", "</fixed>"],
                    num_completions=1,
                    temperature=0,
                ),
            )

            text = pred.completion_text.strip()
            log_probs = pred.completion_logprobs

            print(text)
            print(log_probs)

            probs = [math.exp(log_prob) for log_prob in log_probs]
            print(probs)
            average = sum(probs) / len(log_probs)
            print(average)
            total = math.exp(sum(log_probs))
            print(total)

            print("Total_Probability: ", total)
            print("Average_Probability: ", average)
            print("Log_Probabilities: ", log_probs)
            print("Length: ", len(text))

            # reopen the buggy file and replace the buggy line with the fixed line
            with open(buggy_file, "r+", encoding="locale") as file:
                print("Replacing buggy line with fixed line...")
                buggy_file_contents = file.readlines()
                if check_equal_sign(buggy_file_contents[line_number - 1]):
                    buggy_file_contents[line_number - 1] = (
                        buggy_file_contents[line_number - 1].strip().split("=")[0]
                        + "= "
                        + text
                        + "\n"
                    )
                    if re.sub(
                        r"\s",
                        "",
                        buggy_file_contents[line_number - 1],
                    ) == re.sub(r"\s", "", fixed_line):
                        is_exact = 1
                        correct = correct + 1
                else:
                    buggy_file_contents[line_number - 1] = text + "\n"
                    if re.sub(
                        r"\s",
                        "",
                        buggy_file_contents[line_number - 1],
                    ) == re.sub(r"\s", "", fixed_line):
                        is_exact = 1
                        correct = correct + 1
                print(buggy_file_contents[line_number - 1])

                with open(buggy_file, "w", encoding="locale") as file:
                    print("Writing to file...")
                    file.writelines(buggy_file_contents)

            # get passing tests and total tests
            passing_tests, total_tests = test_passing_data(data[i])

            print("########")

            error_count = 0
            i = i + 1

            ff = open(args.model + "_result_log_prob.txt", "a", encoding="locale")
            ff.write(
                buggy_line.strip()
                + "\t\t"
                + fixed_line.strip()
                + "\t\t"
                + text.strip()
                + "\t\t"
                + str(is_exact)
                + "\t\t"
                + str(total)
                + "\t\t"
                + str(average)
                + "\t\t"
                + str(log_probs)
                + "\t\t"
                + str(len(text))
                + "\t\t"
                + str(total_tests)
                + "\t\t"
                + str(passing_tests)
                + "\n",
            )
            ff.close()

        except Exception as e:
            print(e)
            error_count = error_count = error_count + 1
            if error_count == 12:
                i = i + 1
                error_count = 0
                ff = open(args.model + "_result_log_prob.txt", "a", encoding="locale")
                ff.write("\n")
                ff.close()

            sleep(5)

            continue

    # to dump the prompts used to generate the results
    with open(args.model + "_result_log_prob.jsonl", "w", encoding="locale") as f:
        json.dump(json_objects, f, indent=4)


main()
