from pydantic import BaseModel
from openai import OpenAI
import time

class ExecutionPlan(BaseModel):
    steps: list[str]

class AIAgent:
    def __init__(self, debug=False):
        self.client = OpenAI()
        self.debug = debug

    def landing_page(self, **kwargs):
        prompt = kwargs['prompt']
        goal = "build a modern landing page"
        execution_plan = self.build_execution_plan(goal, prompt)
        messages = [
            {"role": "system", "content": "You must build a landing page for this user based on their requirement."},
            {"role": "user", "content": prompt},
        ]
        for step in execution_plan.steps:
            if self.debug:
                print(f"Executing step: {step}.")
            messages.append({"role": "system", "content": step})
            response = self.ask_model(messages)
            if self.debug:
                print(f"Done: {response}")
            messages.append({"role": "assistant", "content": response})
        messages.append({"role": "user", "content": "Please respond with the final design for my landing page requirement. Please return only HTML and no extra commentary."})
        landing_page = self.ask_model(messages)
        with open("landing.html", "w+") as f:
            f.write(landing_page)
        return "Done. Please view landing.html in your browser."

    def invoke_action(self, action, **kwargs):
        if not hasattr(self, action):
            return "Sorry, I have no idea what to do with this request!?"
        return getattr(self, action)(**kwargs)

    def ask_model(self, messages, model="gpt-4o-mini", temperature=0.7, response_format=None):
        i = 0
        while i < 3:
            i += 1
            try:
                if response_format is not None:
                    response = self.client.beta.chat.completions.parse(
                        model=model,
                        messages=messages,
                        temperature=temperature,
                        response_format=response_format
                    )
                    return response.choices[0].message.parsed
                else:
                    response = self.client.chat.completions.create(
                        model=model,
                        messages=messages,
                        temperature=temperature
                    )
                    return response.choices[0].message.content
            except Exception as ex:
                print(ex)
                time.sleep(2)

    def intent_router(self, user_prompt) -> str:
        prompt = """
            You must analyze the following user prompt and determine which action below best describes the user's request. Respond only with the action e.g. [email]
            1) [email] - the user wishes to send an email.
            2) [landing_page] - the user wishes to build a landing page.
            3) [book_calendar_date] - the user wishes to book a slot in their calendar.
        """
        intent = self.ask_model([
            {"role": "system", "content": prompt},
            {"role": "user", "content": user_prompt}
        ])
        if self.debug:
            print(f"RAW intent: {intent}")
        return intent.replace("[", "").replace("]", "").strip()

    def build_execution_plan(self, goal, prompt):
        system_prompt = f"""
        Given the current goal:'{goal}' and the user prompt, return a step-by-step execution plan for an Agent to work through, so that it can adequately and efficiently fulfill the users request.
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]
        return self.ask_model(messages, response_format=ExecutionPlan)

    def run(self, prompt):
        intent = self.intent_router(prompt)
        if self.debug:
            print(f"Intent for: '{prompt}' === {intent}")
        return self.invoke_action(intent, prompt=prompt)

def main():
    agent = AIAgent(debug=True)
    result = agent.run("I would like to build a landing page for my plumbing business. Please include images for pixels.com and generate copy that makes sense for industry.")
    print(result)

if __name__ == "__main__":
    main()