import os
import yaml
import sys
import importlib
import argparse

from gpt_agent import *
from prompt import *
from task_info import *
from test_gen_code import setup_task_config, run

# Global variable definitions
SCRIPT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "script")
CONFIGS_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "task_config")


def run_code(task_info, las_error=None, message=None):
    """
    Execute generated code for a robot task based on task information and previous errors.

    Args:
        task_info (dict): Dictionary containing task metadata (name, description, etc.).
        las_error (str, optional): Last recorded error message, if any.
        message (list, optional): Message history for the agent.

    Returns:
        tuple: (success_rate, error_message, error_count, run_records)
    """
    check_num = 50
    if message is None:
        message = []

    print("Running code for task:", task_info['task_name'])

    # Extract task info
    task_name = task_info['task_name']
    task_description = task_info['task_description']

    print("Task Name:", task_name)
    print("Task Description:", task_description)

    task, args = setup_task_config(task_name)

    try:
        # Updated to match the new return values of run()
        success_rate, error_message, error_count, run_records = run(task, args, check_num)
        return success_rate, error_message, error_count, run_records

    except KeyboardInterrupt:
        print("Testing interrupted by user.")
        return 0, "Testing interrupted by user", 20

    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        print(f"An error occurred during testing: {e}\n{error_trace}")
        return 0, f"Error during testing: {e}", 20


def main(task_info_dic):
    """
    Main function to test generated code for a given robot task.

    Args:
        task_info_dic (dict): Dictionary containing task information.
    """
    task_info = now_task_info = task_info_dic
    messages = [{
        "role": "system",
        "content": "You need to generate relevant code for some robot tasks in a robot simulation environment based on the provided API."
    }]
    las_error_message = None

    # Run the test
    success_rate, las_error_message, error_count, run_records = run_code(
        now_task_info, las_error_message, messages
    )

    # Evaluate result
    if success_rate >= 0.5:
        print(f"Successfully generated and executed code for task: {task_info['task_name']}")
    else:
        print(f"Failed to generate or execute code for task: {task_info['task_name']}")
        print("Error message:\n", las_error_message)
        now_task_info["task_description"] = (
            f"Failed to generate code, error message: {las_error_message}, "
            f"error count: {str(error_count)}"
        )
        now_task_info["current_code"] = None

    print("Final Success Rate:", success_rate)


if __name__ == "__main__":
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description='Run generated code for a robot task.')
    parser.add_argument('task_name', type=str)
    now_task = None

    # Get task info from task name string
    try:
        task_name = parser.parse_args().task_name.upper()
        exec(f'now_task = {task_name}')
    except Exception as e:
        raise ValueError("Invalid task name specified.") from e

    # Run main function
    main(now_task)


"""
Usage:
python code_gen/run_code.py task_name
""" 