import os
from dotenv import load_dotenv
from swarms.models import OpenAIChat
from swarms.models.stable_diffusion import StableDiffusion
from swarms.structs import Agent, SequentialWorkflow
import swarms.prompts.education as edu_prompts

# Load environment variables
load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")
stability_api_key = os.getenv("STABILITY_API_KEY")

# Initialize language model
llm = OpenAIChat(openai_api_key=api_key, temperature=0.5, max_tokens=3000)

# Initialize Stable Diffusion
sd_api = StableDiffusion(api_key=stability_api_key)

# User preferences (can be dynamically set in a real application)
user_preferences = {
    "subjects": "Cognitive Architectures",
    "learning_style": "Visual",
    "challenge_level": "Moderate"
}

# Formatted prompts from user preferences
curriculum_prompt = edu_prompts.CURRICULUM_DESIGN_PROMPT.format(**user_preferences)
interactive_prompt = edu_prompts.INTERACTIVE_LEARNING_PROMPT.format(**user_preferences)
sample_prompt = edu_prompts.SAMPLE_TEST_PROMPT.format(**user_preferences)
image_prompt = edu_prompts.IMAGE_GENERATION_PROMPT.format(**user_preferences)

# Initialize agents for different educational tasks
curriculum_agent = Agent(llm=llm, max_loops=1, sop=curriculum_prompt)
interactive_learning_agent = Agent(llm=llm, max_loops=1, sop=interactive_prompt)
sample_lesson_agent = Agent(llm=llm, max_loops=1, sop=sample_prompt)

# Create Sequential Workflow
workflow = SequentialWorkflow(max_loops=1)

# Add tasks to workflow with personalized prompts
workflow.add(curriculum_agent, "Generate a curriculum")
workflow.add(interactive_learning_agent, "Generate an interactive lesson")
workflow.add(sample_lesson_agent, "Generate a practice test")

# Execute the workflow for text-based tasks
workflow.run()

# Generate an image using Stable Diffusion
image_result = sd_api.run(image_prompt)

# Output results for each task
for task in workflow.tasks:
    print(f"Task Description: {task.description}\nResult: {task.result}\n")

# Output image result
print(f"Image Generation Task: Generate an image for the interactive lesson\nResult: {image_result}")