from typing import List, Optional
import json
import openai


def call_ai_function(function, args, description, model="gpt-4"):
    # parse args to comma seperated string
    args = ", ".join(args)
    messages = [{"role": "system",
                 "content": f"You are now the following python function: ```# {description}\n{function}```\n\nOnly respond with your `return` value."},
                {"role": "user",
                 "content": args}]

    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0
    )

    return response.choices[0].message["content"]

# Evaluating code


def evaluate_code(code: str) -> List[str]:
    function_string = "def analyze_code(code: str) -> List[str]:"
    args = [code]
    description_string = """Analyzes the given code and returns a list of suggestions for improvements."""

    result_string = call_ai_function(function_string, args, description_string)
    return json.loads(result_string)


# Improving code

def improve_code(suggestions: List[str], code: str) -> str:
    function_string = "def generate_improved_code(suggestions: List[str], code: str) -> str:"
    args = [json.dumps(suggestions), code]
    description_string = """Improves the provided code based on the suggestions provided, making no other changes."""

    result_string = call_ai_function(function_string, args, description_string)
    return result_string


# Writing tests

def write_tests(code: str, focus: List[str]) -> str:
    function_string = "def create_test_cases(code: str, focus: Optional[str] = None) -> str:"
    args = [code, json.dumps(focus)]
    description_string = """Generates test cases for the existing code, focusing on specific areas if required."""

    result_string = call_ai_function(function_string, args, description_string)
    return result_string