import random
import os
from dotenv import load_dotenv
from swarms.models import OpenAIChat
from swarms.structs import Agent
from swarms.models.stable_diffusion import StableDiffusion

load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY")
stability_api_key = os.getenv("STABILITY_API_KEY")

# Initialize the language model and image generation model
llm = OpenAIChat(openai_api_key=openai_api_key, temperature=0.5, max_tokens=3000)
sd_api = StableDiffusion(api_key=stability_api_key)

# Creative Concept Generator for Product Ads
class ProductAdConceptGenerator:
    def __init__(self, product_name):
        self.product_name = product_name
        self.themes = [
            "futuristic", "rustic", "luxurious", "minimalistic", "vibrant", "elegant",
            "retro", "urban", "ethereal", "surreal", "artistic", "tech-savvy",
            "vintage", "natural", "sophisticated", "playful", "dynamic", "serene", "lasers," "lightning"
        ]
        self.contexts = [
            "in an everyday setting", "in a rave setting", "in an abstract environment", 
            "in an adventurous context", "surrounded by nature", "in a high-tech setting",
            "in a historical context", "in a busy urban scene", "in a tranquil and peaceful setting",
            "against a backdrop of city lights", "in a surreal dreamscape", "in a festive atmosphere",
            "in a luxurious setting", "in a playful and colorful background", "in an ice cave setting",
            "in a serene and calm landscape"
        ]

    def generate_concept(self):
        theme = random.choice(self.themes)
        context = random.choice(self.contexts)
        return f"An ad for {self.product_name} that embodies a {theme} theme {context}"

# User input
product_name = input("Enter a product name for ad creation (e.g., 'PS5', 'AirPods', 'Kirkland Vodka'): ")
social_media_platform = input("Enter a social media platform (e.g., 'Facebook', 'Twitter', 'Instagram'): ")

# Generate creative concept
concept_generator = ProductAdConceptGenerator(product_name)
creative_concept = concept_generator.generate_concept()

# Generate product image based on the creative concept
image_paths = sd_api.run(creative_concept)

# Generate ad copy
ad_copy_agent = Agent(llm=llm, max_loops=1)
ad_copy_prompt = f"Write a compelling {social_media_platform} ad copy for a product photo showing {product_name} {creative_concept}."
ad_copy = ad_copy_agent.run(task=ad_copy_prompt)

# Output the results
print("Creative Concept:", creative_concept)
print("Image Path:", image_paths[0] if image_paths else "No image generated")
print("Ad Copy:", ad_copy)