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