import sys
import os

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from gpt_api.gpt_agent import *
from gpt_api.prompt import *
from gpt_api.task_info import *
from gpt_api.task_info import get_all_tasks
from test_gpt_code import *  

import argparse

def generate_code(task_info, las_error=None, message=None):
    """Generate code for robot task based on task info and previous errors."""
    if message is None:
        message = []
        
    # Extract task information
    task_name = task_info['task_name']
    task_description = task_info['task_description']
    current_code = task_info['current_code']
    available_env_function = str(AVAILABLE_ENV_FUNCTOIN)
    function_example = str(FUNCTION_EXAMPLE)
    available_constants = str(AVAILABLE_CONSTANTS)

    # Generate code based on error status
    if las_error is not None:
        # Handle error case - provide error info to improve generation
        actor_name_keys, actor_data_keys, actor_points_discription = get_actor_keys_and_points_discription(f"gpt_{task_name}")
        Prompt = (
            f"The code is unsuccessful, \nLast Error Message: \n{las_error}\n\n"
            f"Task Discription: \n{task_description}\n\n"
            f"The Actor Points Discription: {actor_points_discription}"
        )
    else:
        # First attempt case - create initial code file
        res = f'''
from .base_task import Base_task
from .{task_name} import {task_name}
from .utils import *
import sapien

class gpt_{task_name}({task_name}):
    def play_once(self):
        pass
        '''
        file_name = f"envs/gpt_{task_name}.py"
        with open(file_name, 'w') as file:
            file.write(res)
            
        # Get actor information for prompt
        actor_name_keys, actor_data_keys, actor_points_discription = get_actor_keys_and_points_discription(f"gpt_{task_name}")
        
        # Construct full prompt with all necessary information
        Prompt = (
            f"{BASIC_INFO}\n\n"
            f"Task Discription: \n{task_description}\n\n"
            f"Available API: \n{available_env_function}\n\n"
            f"Function Example: \n{function_example}\n\n"
            f"Available Constants: \n{available_constants}\n\n"
            f"The Actor Name List: {actor_name_keys}\n\n"
            f"The Actor Data List: {actor_data_keys}\n\n"
            f"The Actor Points Discription: {actor_points_discription}\n\n"
            f"Current Code:\n{current_code}"
        )
    
    # Add prompt to message history
    message.append({"role": "user", "content": Prompt})

    # Generate code using the model
    res = generate(message, gpt="deepseek", temperature=0.5)
    
    # Extract the relevant portion of the generated code
    res = f'''
from .base_task import Base_task
from .{task_name} import {task_name}
from .utils import *
import sapien

class gpt_{task_name}({task_name}):
    ''' + res[res.find('def play_once'):res.rfind("```")]
    
    # Save generated code to file
    file_name = f"envs/gpt_{task_name}.py"
    with open(file_name, 'w') as file:
        file.write(res)
    
    # Test the generated code
    success_rate, error_message, error_count = test_run(f"gpt_{task_name}")
    
    return res, success_rate, error_message, error_count