import os
import json
from langchain_openai import AzureChatOpenAI
from langchain.schema import HumanMessage, SystemMessage
import time
import argparse


def read_prompt_from_file(filename):
    with open(filename, "r") as file:
        prompt = file.read()
    return prompt


def getAGPT4Instance(
    depoyment_name="gpt-4",
    temperature=0.7,
    request_timeout=60,
    api_version="2023-05-15",
    top_p=0,
):
    gpt4_api_key = os.getenv("OPENAI_API_KEY_GPT4")
    gpt4_api_depoyment_name = (
        depoyment_name 
    )
    gpt4_api_endpoint = os.getenv("OPENAI_API_ENDPOINT_GPT4")

    return AzureChatOpenAI(
        deployment_name=gpt4_api_depoyment_name,
        api_key=gpt4_api_key,
        api_version=api_version,
        azure_endpoint=gpt4_api_endpoint,
        temperature=temperature,
        request_timeout=request_timeout,
        response_format={"type": "json_object"},
    )


#


def getParsedContent(llm, template, elementNames=None, systemMsg=None):
    counter = 0
    if elementNames == None:
        if systemMsg != None:
            flag, output = callLLM(
                llm, [SystemMessage(content=systemMsg), HumanMessage(content=template)]
            )

        else:
            flag, output = callLLM(llm, [HumanMessage(content=template)])
        return flag, output.content

    while counter < 5:
        counter = counter + 1
        try:

            flag, output = callLLM(
                llm, [HumanMessage(content=template)]
            )  # llm([HumanMessage(content=template)])
            if not flag:
                print("Error: %s" % output)
                return flag, output

            parsed_output = json.loads(output.content)
            parsed_elements = {}

            for elementName in elementNames:
                parsedElement = parsed_output.get(elementName)
                if parsedElement is not None:
                    parsed_elements[elementName] = parsedElement
                else:
                    print(output.content)
                    print(f"Error: '{elementName}' key not found in the output")
                    break
            else:
                return True, parsed_elements

        except json.JSONDecodeError as e:
            print(output.content)
            print("Error: Incorrect JSON format", str(e))
            continue

    return False, "Error: Incorrect JSON format"


def callLLM(llm, inputObject):
    try:
        output = llm.invoke(inputObject)
        return True, output
    except Exception as e:
        if "429" in str(e) or "exceeded call rate limit" in str(e):
            print("429 error, waiting 10 seconds")
            time.sleep(1)
            return callLLM(llm, inputObject)
        return False, str((e))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Process JSON file to add result fields."
    )
    parser.add_argument(
        "--input", type=str, required=True, help="Path to the input JSON file."
    )
    parser.add_argument(
        "--output", type=str, required=True, help="Path to the output JSON file."
    )
    parser.add_argument("--model", type=str, default="gpt-4o", help="Model to use.")
    parser.add_argument(
        "--prompt",
        type=str,
        default="../prompts/judgePromptStrict.txt",
        help="Prompt file to use.",
    )
    parser.add_argument(
        "--start_idx", type=int, default=0, help="start_idx in input file."
    )
    parser.add_argument(
        "--end_idx", type=int, default=-1, help="end_idx in input file."
    )

    args = parser.parse_args()
    model_name = args.model
    llmInstance = getAGPT4Instance(
        depoyment_name=model_name, api_version="2024-09-01-preview"
    )
    judgeInput = read_prompt_from_file(args.prompt)
    with open(args.input, "r") as input_file:
        data = json.load(input_file)
        data = data[args.start_idx : args.end_idx]

    for idx, entry in enumerate(data):
        try:
            if idx % 10 == 0:
                print("Processing entry %d/%d" % (idx, len(data)))

            if "real" in entry:
                prompt_key = "real"
            elif "hypothetical" in entry:
                prompt_key = "hypothetical"
            else:
                prompt_key = "prompt"
            task = entry[prompt_key]
            steered_reasoning = entry["steered_result"].split("</think>")[0]
            steered_answer = entry["steered_result"].split("</think>")[1]
            filledTemplate = judgeInput.format(
                task=task, reasoning=steered_reasoning, answer=steered_answer
            )
            flag, parsedOutput = getParsedContent(
                llmInstance,
                filledTemplate,
                elementNames=[
                    "model_awareness",
                    "hypothetical_recognition",
                    "task_performance",
                ],
            )
            for i in range(5):
                if not flag:
                    print("Error in parsing the REAL output")
                    print("Error: %s" % parsedOutput)
                    print(i)
                    entry["judge_steered_result"] = {}
                    entry["judge_steered_result"]["steered_model_awareness"] = ""
                    entry["judge_steered_result"]["steered_recognition"] = ""
                    entry["judge_steered_result"]["steered_task_performance"] = ""
                    continue
                else:
                    entry["judge_steered_result"] = {}
                    entry["judge_steered_result"]["steered_model_awareness"] = (
                        parsedOutput["model_awareness"]
                    )
                    entry["judge_steered_result"]["steered_recognition"] = parsedOutput[
                        "hypothetical_recognition"
                    ]
                    entry["judge_steered_result"]["steered_task_performance"] = (
                        parsedOutput["task_performance"]
                    )
                    break
            # Save the modified data to the output file
            with open(args.output, "w") as output_file:
                json.dump(data, output_file, indent=4)
        except:
            continue
