import concurrent.futures
from typing import Dict, Optional
import secrets
import string
import uuid

from dotenv import load_dotenv
from swarms import Agent

import replicate

from swarms.utils.str_to_dict import str_to_dict

load_dotenv()


def generate_key(prefix: str = "run") -> str:
    """
    Generates an API key similar to OpenAI's format (sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX).

    Args:
        prefix (str): The prefix for the API key. Defaults to "sk".

    Returns:
        str: An API key string in format: prefix-<48 random characters>
    """
    # Create random string of letters and numbers
    alphabet = string.ascii_letters + string.digits
    random_part = "".join(secrets.choice(alphabet) for _ in range(28))
    return f"{prefix}-{random_part}"


def _generate_media(
    prompt: str = None, modalities: list = None
) -> Dict[str, str]:
    """
    Generate media content (images or videos) based on text prompts using AI models.

    Args:
        prompt (str): Text description of the content to be generated
        modalities (list): List of media types to generate (e.g., ["image", "video"])

    Returns:
        Dict[str, str]: Dictionary containing file paths of generated media
    """
    if not prompt or not modalities:
        raise ValueError("Prompt and modalities must be provided")

    input = {"prompt": prompt}
    results = {}

    def _generate_image(input: Dict) -> str:
        """Generate an image and return the file path."""
        output = replicate.run(
            "black-forest-labs/flux-dev", input=input
        )
        file_paths = []

        for index, item in enumerate(output):
            unique_id = str(uuid.uuid4())
            artifact = item.read()
            file_path = f"output_{unique_id}_{index}.webp"

            with open(file_path, "wb") as file:
                file.write(artifact)

            file_paths.append(file_path)

        return file_paths

    def _generate_video(input: Dict) -> str:
        """Generate a video and return the file path."""
        output = replicate.run("luma/ray", input=input)
        unique_id = str(uuid.uuid4())
        artifact = output.read()
        file_path = f"output_{unique_id}.mp4"

        with open(file_path, "wb") as file:
            file.write(artifact)

        return file_path

    for modality in modalities:
        if modality == "image":
            results["images"] = _generate_image(input)
        elif modality == "video":
            results["video"] = _generate_video(input)
        else:
            raise ValueError(f"Unsupported modality: {modality}")

    print(results)

    return results


def generate_media(
    modalities: list,
    prompt: Optional[str] = None,
    count: int = 1,
) -> Dict:
    with concurrent.futures.ThreadPoolExecutor(
        max_workers=count
    ) as executor:
        # Create list of identical tasks to run concurrently
        futures = [
            executor.submit(
                _generate_media,
                prompt=prompt,  # Fix: Pass as keyword arguments
                modalities=modalities,
            )
            for _ in range(count)
        ]

        # Wait for all tasks to complete and collect results
        results = [
            future.result()
            for future in concurrent.futures.as_completed(futures)
        ]

    return {"results": results}


tools = [
    {
        "type": "function",
        "function": {
            "name": "generate_media",
            "description": "Generate different types of media content (image, video, or music) based on text prompts using AI models.",
            "parameters": {
                "type": "object",
                "properties": {
                    "modality": {
                        "type": "array",
                        "items": {
                            "type": "string",
                            "enum": ["image", "video", "music"],
                        },
                        "description": "The type of media content to generate",
                    },
                    "prompt": {
                        "type": "string",
                        "description": "Text description of the content to be generated. Specialize it for the modality at hand. For example, if you are generating an image, the prompt should be a description of the image you want to see. If you are generating a video, the prompt should be a description of the video you want to see. If you are generating music, the prompt should be a description of the music you want to hear.",
                    },
                    "count": {
                        "type": "integer",
                        "description": "Number of outputs to generate (1-4)",
                    },
                },
                "required": [
                    "modality",
                    "prompt",
                    "count",
                ],
            },
        },
    }
]


MEDIA_GENERATION_SYSTEM_PROMPT = """
You are an expert AI Media Generation Assistant, specialized in crafting precise and effective prompts for generating images, videos, and music. Your role is to help users create high-quality media content by understanding their requests and translating them into optimal prompts.

GENERAL GUIDELINES:
- Always analyze the user's request carefully to determine the appropriate modality (image, video, or music)
- Maintain a balanced level of detail in prompts - specific enough to capture the desired outcome but not overly verbose
- Consider the technical limitations and capabilities of AI generation systems
- When unclear, ask for clarification about specific details or preferences

MODALITY-SPECIFIC GUIDELINES:

1. IMAGE GENERATION:
- Structure prompts with primary subject first, followed by style, mood, and technical specifications
- Include relevant art styles when specified (e.g., "digital art", "oil painting", "watercolor", "photorealistic")
- Consider composition elements (foreground, background, lighting, perspective)
- Use specific adjectives for clarity (instead of "beautiful", specify "vibrant", "ethereal", "gritty", etc.)

Example image prompts:
- "A serene Japanese garden at sunset, with cherry blossoms falling, painted in traditional ukiyo-e style, soft pastel colors"
- "Cyberpunk cityscape at night, neon lights reflecting in rain puddles, hyper-realistic digital art style"

2. VIDEO GENERATION:
- Describe the sequence of events clearly
- Specify camera movements if relevant (pan, zoom, tracking shot)
- Include timing and transitions when necessary
- Focus on dynamic elements and motion

Example video prompts:
- "Timelapse of a flower blooming in a garden, close-up shot, soft natural lighting, 10-second duration"
- "Drone shot flying through autumn forest, camera slowly rising above the canopy, revealing mountains in the distance"

3. MUSIC GENERATION:
- Specify genre, tempo, and mood
- Mention key instruments or sounds
- Include emotional qualities and intensity
- Reference similar artists or styles if relevant

Example music prompts:
- "Calm ambient electronic music with soft synthesizer pads, gentle piano melodies, 80 BPM, suitable for meditation"
- "Upbeat jazz fusion track with prominent bass line, dynamic drums, and horn section, inspired by Weather Report"

COUNT HANDLING:
- When multiple outputs are requested (1-4), maintain consistency while introducing subtle variations
- For images: Vary composition or perspective while maintaining style
- For videos: Adjust camera angles or timing while keeping the core concept
- For music: Modify instrument arrangements or tempo while preserving the genre and mood

PROMPT OPTIMIZATION PROCESS:
1. Identify core requirements from user input
2. Determine appropriate modality
3. Add necessary style and technical specifications
4. Adjust detail level based on complexity
5. Consider count and create variations if needed

EXAMPLES OF HANDLING USER REQUESTS:

User: "I want a fantasy landscape"
Assistant response: {
    "modality": "image",
    "prompt": "Majestic fantasy landscape with floating islands, crystal waterfalls, and ancient magical ruins, ethereal lighting, digital art style with rich colors",
    "count": 1
}

User: "Create 3 variations of a peaceful nature scene"
Assistant response: {
    "modality": "image",
    "prompt": "Tranquil forest clearing with morning mist, sunbeams filtering through ancient trees, photorealistic style with soft natural lighting",
    "count": 1
}

IMPORTANT CONSIDERATIONS:
- Avoid harmful, unethical, or inappropriate content
- Respect copyright and intellectual property guidelines
- Maintain consistency with brand guidelines when specified
- Consider technical limitations of current AI generation systems

"""

# Initialize the agent with the new system prompt
agent = Agent(
    agent_name="Media-Generation-Agent",
    agent_description="AI Media Generation Assistant",
    system_prompt=MEDIA_GENERATION_SYSTEM_PROMPT,
    max_loops=1,
    tools_list_dictionary=tools,
    output_type="final",
)


def create_agent(task: str):
    output = str_to_dict(agent.run(task))

    print(output)
    print(type(output))

    prompt = output["prompt"]
    count = output["count"]
    modalities = output["modality"]

    output = generate_media(
        modalities=modalities,
        prompt=prompt,
        count=count,
    )

    run_id = generate_key()

    total_cost = 0

    for modality in modalities:
        if modality == "image":
            total_cost += 0.1
        elif modality == "video":
            total_cost += 1

    result = {
        "id": run_id,
        "success": True,
        "prompt": prompt,
        "count": count,
        "modality": modalities,
        "total_cost": total_cost,
    }

    return result


if __name__ == "__main__":
    task = "Create 3 super kawaii variations of a magical Chinese mountain garden scene in anime style! 🌸✨ Include adorable elements like: cute koi fish swimming in crystal ponds, fluffy clouds floating around misty peaks, tiny pagodas with twinkling lights, and playful pandas hiding in bamboo groves. Make it extra magical with sparkles and soft pastel colors! Create both a video and an image for each variation. Just 1."
    output = create_agent(task)
    print("✨ Yay! Here's your super cute creation! ✨")
    print(output)