import sys
import openai
import autogen
import time
import os
from config import openai_config
from openai import AzureOpenAI
import traceback

def run_code(cell):
    """
    Returns the path to the python interpreter.
    """
    # import prompts
    from prompts_mimic import CodeHeader
    try:
        global_var = {"answer": 0}
        exec(CodeHeader+cell, global_var)
        cell = "\n".join([line for line in cell.split("\n") if line.strip() and not line.strip().startswith("#")])
        if not 'answer' in cell.split('\n')[-1]:
            return "Please save the answer to the question in the variable 'answer'."
        return str(global_var['answer'])
    except Exception as e:
        error_info = traceback.format_exc()
        code = CodeHeader + cell
        if "SyntaxError" in str(repr(e)):
            error_line = str(repr(e))
            
            error_type = error_line.split('(')[0]
            # then parse out the error message
            error_message = error_line.split(',')[0].split('(')[1]
            # then parse out the error line
            error_line = error_line.split('"')[1]
        elif "KeyError" in str(repr(e)):
            code = code.split('\n')
            key = str(repr(e)).split("'")[1]
            error_type = str(repr(e)).split('(')[0]
            for i in range(len(code)):
                if key in code[i]:
                    error_line = code[i]
            error_message = str(repr(e))
        elif "TypeError" in str(repr(e)):
            error_type = str(repr(e)).split('(')[0]
            error_message = str(e)
            function_mapping_dict = {"get_value": "GetValue", "data_filter": "FilterDB", "db_loader": "LoadDB", "sql_interpreter": "SQLInterpreter", "date_calculator": "Calendar"}
            error_key = ""
            for key in function_mapping_dict.keys():
                if key in error_message:
                    error_message = error_message.replace(key, function_mapping_dict[key])
                    error_key = function_mapping_dict[key]
            code = code.split('\n')
            error_line = ""
            for i in range(len(code)):
                if error_key in code[i]:
                    error_line = code[i]
        else:
            error_type = ""
            error_message = str(repr(e)).split("('")[-1].split("')")[0]
            error_line = ""
        # use one sentence to introduce the previous parsed error information
        if error_type != "" and error_line != "":
            error_info = f'{error_type}: {error_message}. The error messages occur in the code line "{error_line}".'
        else:
            error_info = f'Error: {error_message}.'
        error_info += '\nPlease make modifications accordingly and make sure the rest code works well with the modification.'

        return error_info


def llm_agent(config_list):
    llm_config = {
        "functions": [
            {
                "name": "python",
                "description": "run cell in ipython and return the execution result.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "cell": {
                            "type": "string",
                            "description": "Valid Python cell to execute.",
                        }
                    },
                    "required": ["cell"],
                },
            },
        ],
        "config_list": config_list,
        "request_timeout": 120,
    }
    chatbot = autogen.AssistantAgent(
        name="chatbot",
        system_message="For coding tasks, only use the functions you have been provided with. Reply TERMINATE when the task is done.",
        llm_config=llm_config,
    )
    return chatbot