import os
import time
import requests
from PIL import Image
from io import BytesIO

from src.menu_design.core.menu_design import STATIC_PATH
from src.utils.settings import settings

# Environment variables
IDEOGRAM_API_KEY = settings.IDEOGRAM_API_KEY
IDEOGRAM_URL = settings.IDEOGRAM_URL

# Base URL used to construct image links
BASE_URL = settings.BASE_URL
#BASE_URL = "http://127.0.0.1:8000/"

def generate_image(desc: str, layout: str, style: str, image_number: int = 1):
    """
    Generate and save multiple images based on description, layout, and style using Ideogram API.
    
    Args:
        desc (str): Description of the content.
        layout (str): Layout keyword to include in the prompt.
        style (str): Style keyword to include in the prompt.
        image_number (int): Number of images to generate (default is 1).
    
    Returns:
        List[str]: URLs of the saved images.
    """
    image_urls = []

    for i in range(image_number):
        try:
            # Step 1: Construct the image prompt
            image_prompt = f"{desc}. Style: {style}, layout: {layout}"
            print(f"Generating image with prompt: {image_prompt}")

            # Step 2: Call Ideogram API to generate the image
            response = requests.post(
                IDEOGRAM_URL,
                headers={
                    "Api-Key": IDEOGRAM_API_KEY,
                },
                files={
                    "prompt": (None, image_prompt),
                    "aspect_ratio": (None, "1x1"),
                }
            )

            print(f"Received response from Ideogram API: {response.status_code}")

            # Step 3: Handle errors
            if response.status_code != 200:
                raise Exception(f"Ideogram API error: {response.text}")

            # Step 4: Extract image URL from the API response
            generation_url = response.json().get('data')[0].get('url')
            if not generation_url:
                raise Exception("No image URL returned by Ideogram API.")

            # Step 5: Download the generated image
            image_response = requests.get(generation_url)
            if image_response.status_code != 200:
                raise Exception("Failed to download the generated image.")

            # Step 6: Save image locally
            timestamp = int(time.time())
            image_filename = f"image_{i+1}_{timestamp}.png"
            image_filepath = os.path.join(STATIC_PATH, image_filename)

            img = Image.open(BytesIO(image_response.content))
            img.save(image_filepath)
            print(f"Saved image at: {image_filepath}")

            # Step 7: Create public URL
            full_image_url = f"{BASE_URL}static/generated_images/{image_filename}"
            image_urls.append(full_image_url)

        except Exception as e:
            print(f"Error generating image {i + 1}: {str(e)}")

    return image_urls