import click
from airtrain.integrations.openai.skills import OpenAIChatSkill, OpenAIInput
import os
from dotenv import load_dotenv

load_dotenv()


def initialize_chat():
    return OpenAIChatSkill()


@click.group()
def cli():
    """Airtrain CLI - Your AI Agent Building Assistant"""
    pass


@cli.command()
def chat():
    """Start an interactive chat session with Airtrain"""
    skill = initialize_chat()
    click.echo("Welcome to Airtrain! I'm here to help you build your AI Agent.")
    click.echo("Type 'exit' to end the conversation.\n")

    while True:
        user_input = click.prompt("You", type=str)

        if user_input.lower() == "exit":
            click.echo("\nGoodbye! Have a great day!")
            break

        try:
            input_data = OpenAIInput(
                user_input=user_input,
                system_prompt="You are an AI assistant that helps users build their own AI agents. Be helpful and provide clear explanations.",
                model="gpt-4o",
                temperature=0.7,
            )

            result = skill.process(input_data)
            click.echo(f"\nAirtrain: {result.response}\n")

        except Exception as e:
            click.echo(f"\nError: {str(e)}\n")


def main():
    cli()


if __name__ == "__main__":
    main()
