import os
import json
from langchain_openai import AzureChatOpenAI
from langchain.schema import HumanMessage, SystemMessage
import time
import argparse
from typing import List, Dict, Any, Tuple
import concurrent.futures


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 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(10)  # Increased wait time for rate limits
            return callLLM(llm, inputObject)
        return False, str(e)


def callLLMBatch(llm, batch_inputs):
    """Process a batch of inputs in parallel using concurrent.futures"""
    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        # Submit all tasks
        future_to_input = {
            executor.submit(llm.invoke, input_obj): i
            for i, input_obj in enumerate(batch_inputs)
        }

        # Process results as they complete
        for future in concurrent.futures.as_completed(future_to_input):
            input_idx = future_to_input[future]
            try:
                result = future.result()
                results.append((input_idx, True, result))
            except Exception as e:
                if "429" in str(e) or "exceeded call rate limit" in str(e):
                    # For rate limit errors, retry the specific input
                    print(f"429 error for input {input_idx}, retrying after wait")
                    time.sleep(10)
                    try:
                        result = llm.invoke(batch_inputs[input_idx])
                        results.append((input_idx, True, result))
                    except Exception as retry_e:
                        results.append((input_idx, False, str(retry_e)))
                else:
                    results.append((input_idx, False, str(e)))

    # Sort results by original index
    results.sort(key=lambda x: x[0])
    return [(success, output) for _, success, output in results]


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)])
            if not flag:
                print(f"Error: {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 process_batch(llm, batch_templates, elementNames):
    """Process a batch of templates and parse their outputs"""
    # Create message objects for each template
    batch_inputs = [[HumanMessage(content=template)] for template in batch_templates]

    # Call LLM with batch inputs
    batch_results = callLLMBatch(llm, batch_inputs)

    parsed_results = []
    for i, (flag, output) in enumerate(batch_results):
        if not flag:
            parsed_results.append((False, output))
            continue

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

            all_keys_found = True
            for elementName in elementNames:
                parsedElement = parsed_output.get(elementName)
                if parsedElement is not None:
                    parsed_elements[elementName] = parsedElement
                else:
                    print(f"Template {i}: '{elementName}' key not found in output")
                    all_keys_found = False
                    break

            if all_keys_found:
                parsed_results.append((True, parsed_elements))
            else:
                parsed_results.append((False, "Missing required keys in output"))

        except json.JSONDecodeError as e:
            print(f"Template {i}: Incorrect JSON format - {str(e)}")
            parsed_results.append((False, f"Error: Incorrect JSON format - {str(e)}"))

    return parsed_results


def prepare_data_for_batching(data, judge_input, batch_size=10):
    """Prepare data in batches for processing"""
    real_batch_templates = []
    hypothetical_batch_templates = []

    for entry in data:
        if not (
            "</think>" in entry["hypothetical_result"]
            and "</think>" in entry["real_result"]
        ):
            continue

        # Prepare real data template
        real_task = entry["real"]
        real_reasoning = entry["real_result"].split("</think>")[0]
        real_answer = entry["real_result"].split("</think>")[1]
        real_template = judge_input.format(
            task=real_task, reasoning=real_reasoning, answer=real_answer
        )
        real_batch_templates.append(real_template)

        # Prepare hypothetical data template
        hypothetical_task = entry["hypothetical"]
        hypothetical_reasoning = entry["hypothetical_result"].split("</think>")[0]
        hypothetical_answer = entry["hypothetical_result"].split("</think>")[1]
        hypothetical_template = judge_input.format(
            task=hypothetical_task,
            reasoning=hypothetical_reasoning,
            answer=hypothetical_answer,
        )
        hypothetical_batch_templates.append(hypothetical_template)

    # Group into batches
    real_batches = [
        real_batch_templates[i : i + batch_size]
        for i in range(0, len(real_batch_templates), batch_size)
    ]
    hypothetical_batches = [
        hypothetical_batch_templates[i : i + batch_size]
        for i in range(0, len(hypothetical_batch_templates), batch_size)
    ]

    return real_batches, hypothetical_batches


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."
    )
    parser.add_argument(
        "--batch_size", type=int, default=10, help="Batch size for API calls."
    )

    args = parser.parse_args()
    model_name = args.model
    batch_size = args.batch_size
    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]

    # Filter and prepare valid entries
    valid_entries = [
        entry
        for entry in data
        if (
            "</think>" in entry.get("hypothetical_result", "")
            and "</think>" in entry.get("real_result", "")
        )
    ]

    # Prepare batches
    real_batches, hypothetical_batches = prepare_data_for_batching(
        valid_entries, judgeInput, batch_size=batch_size
    )

    # Process batches
    element_names = ["model_awareness", "hypothetical_recognition", "task_performance"]

    # Process real entries
    print(f"Processing {len(valid_entries)} entries in batches of {batch_size}")
    valid_entry_idx = 0

    for batch_idx, (real_batch, hypothetical_batch) in enumerate(
        zip(real_batches, hypothetical_batches)
    ):
        print(f"Processing batch {batch_idx+1}/{len(real_batches)}")

        # Process real entries batch
        real_results = process_batch(llmInstance, real_batch, element_names)

        # Process hypothetical entries batch
        hypothetical_results = process_batch(
            llmInstance, hypothetical_batch, element_names
        )

        # Update entries with results
        for i in range(len(real_batch)):
            if valid_entry_idx >= len(valid_entries):
                break

            entry = valid_entries[valid_entry_idx]
            data_idx = data.index(entry)  # Find index in original data

            # Update with real results
            real_flag, real_output = real_results[i]
            if real_flag:
                data[data_idx]["real_model_awareness"] = real_output["model_awareness"]
                data[data_idx]["real_recognition"] = real_output[
                    "hypothetical_recognition"
                ]
                data[data_idx]["real_task_performance"] = real_output[
                    "task_performance"
                ]
            else:
                print(
                    f"Error in parsing real output for entry {valid_entry_idx}: {real_output}"
                )
                data[data_idx]["real_model_awareness"] = ""
                data[data_idx]["real_recognition"] = ""
                data[data_idx]["real_task_performance"] = ""

            # Update with hypothetical results
            hypo_flag, hypo_output = hypothetical_results[i]
            if hypo_flag:
                data[data_idx]["hypothetical_model_awareness"] = hypo_output[
                    "model_awareness"
                ]
                data[data_idx]["hypothetical_recognition"] = hypo_output[
                    "hypothetical_recognition"
                ]
                data[data_idx]["hypothetical_task_performance"] = hypo_output[
                    "task_performance"
                ]
            else:
                print(
                    f"Error in parsing hypothetical output for entry {valid_entry_idx}: {hypo_output}"
                )
                data[data_idx]["hypothetical_model_awareness"] = ""
                data[data_idx]["hypothetical_recognition"] = ""
                data[data_idx]["hypothetical_task_performance"] = ""

            valid_entry_idx += 1

        # Save intermediate results after each batch
        with open(args.output, "w") as output_file:
            json.dump(data, output_file, indent=4)

    print(f"Processing complete. Results saved to {args.output}")
