from MAIAI import Agent
from openai import OpenAI
import json

client = OpenAI()

class Task:
    def __init__(self, agent: Agent, goal: str, expected_output: str = ""):
        self.agent = agent
        self.goal = goal
        self.expected_output = expected_output
        
    def execute_clean(self):
            try:
                completion = client.chat.completions.create(
                    model=self.agent.model,
                    temperature=self.agent.temperature,
                    messages=[
                        {"role": "system", "content": f"{self.agent.role}"}, 
                        {"role": "user", "content": f"""{self.goal}"""}
                    ]
                )
                output = completion.choices[0].message.content
            except Exception as e:
                output = f"An error occurred: {str(e)}"
            return output
    
    def execute(self, response_type = ""):

        # Response in JSON
        if response_type == "json":
            completion = client.chat.completions.create(
            model=self.agent.model,
            temperature=self.agent.temperature,
            response_format={ "type": "json_object" },
            messages=[
                {"role": "system", "content": f"{self.agent.role}"}, 
                {"role": "user", "content": f"""You MUST give your response as json in following format: {self.expected_output}. 
                 {self.goal}"""}
                ]
            )

            try:
                output = completion.choices[0].message.content
                output = json.loads(output)

            except Exception as e:
                output = "Unable to load JSON file"
        
        # Normal response with no response type requirements
        else:    
            completion = client.chat.completions.create(
                model=self.agent.model,
                temperature=self.agent.temperature,
                messages=[
                    {"role": "system", "content": f"{self.agent.role}"}, 
                    {"role": "user", "content": f"""You MUST give your response as {self.expected_output}. 
                    {self.goal}"""}
                ]
            )

            try:
                output = completion.choices[0].message.content
            except Exception as e:
                output = None

        return output