#!/usr/bin/env python

import os
import argparse
from dotenv import load_dotenv
from openai import OpenAI

def load_openai_key():
    # Load environment variables in a file called .env
    load_dotenv(override=True)
    api_key = os.getenv('OPENAI_API_KEY')

    # Check the key
    if not api_key:
        return "Error: No API key was found!"
    elif not api_key.startswith("sk-proj-"):
        return "Error: An API key was found, but it doesn't start sk-proj-; please check you're using the right key"
    elif api_key.strip() != api_key:
        return "Error: An API key was found, but it looks like it might have space or tab characters at the start or end - please remove them!"
    else:
        return "API key found and looks good so far!"

def ask_llm(client, model, user_prompt):
    system_prompt = """
    you are a writing assistant with an expertise in children's stories.
    Write a bedtime story inspired by the subject below.
    The story should have a begining, middle, and end.
    The story shoukd be appropriate for children ages 5-8 and have a positive message.
    I should be able to read the entire story in about 3 minutes
    """
    response = client.chat.completions.create(
            model = model,
            messages = [ {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}]
            )
    return response.choices[0].message.content

def main():
    parser = argparse.ArgumentParser(description="AI Bedtime Storyteller")
    parser.add_argument("provider", choices=["openai", "ollama"], help="AI provider to use")
    parser.add_argument("--model", help="Model to use for Ollama (required if provider is 'ollama')", required="ollama" in parser.parse_known_args()[0].provider)
    parser.add_argument("subject", help="What do you want the story to be about?")

    args = parser.parse_args()

    if args.provider == "openai":
        load_openai_key()
        client = OpenAI()
        model = "gpt-4o-mini"
    elif args.provider == "ollama":
        client = OpenAI(base_url='http://localhost:11434/v1', api_key='ollama')
        model = args.model
    else:
        return "Error: invalid provider!"

    user_prompt = args.subject

    result = ask_llm(client, model, user_prompt)
    print("AI Response:", result)

if __name__ == "__main__":
    main()