import json

def create_openai_prompt_format(json_file_path):
    """
    Reads assignment problems from a JSON file and converts each into
    an OpenAI-style list of messages.
    
    :param json_file_path: Path to the 'assignment_problems.json' file
    :return: A list of conversations (each conversation is a list of messages).
    """
    with open(json_file_path, 'r') as f:
        problems = json.load(f)

    all_conversations = []

    system_message = {
        "role": "system",
        "content": (
            "You are ChatGPT, a large language model trained by OpenAI. "
            "You will be given a bipartite graph (assignment problem). "
            "Your goal is to find the minimum total cost matching using the Hungarian algorithm, "
            "and return both the minimal cost and the matching pairs."
        )
    }

    for problem in problems:
        user_message = {
            "role": "user",
            "content": (
                f"Problem ID: {problem['problem_id']}\n"
                f"Bipartite Graph:\n{json.dumps(problem['graph'], indent=2)}\n\n"
                "Please find the minimal total cost assignment and the matching pairs."
            )
        }
        
        conversation = [system_message, user_message]
        all_conversations.append(conversation)

    return all_conversations


if __name__ == "__main__":
    json_file_path = "assignment_problems.json"
    
    conversations = create_openai_prompt_format(json_file_path)
    
    with open("assignment_prompts.json", "w") as outfile:
        json.dump(conversations, outfile, indent=2)
    
    print("Generated assignment_prompts.json with OpenAI-formatted messages.")
