import streamlit as st
import requests
import os
import base64
from PIL import Image
from io import BytesIO

# Set page title and layout
st.set_page_config(page_title="AI Image Generator", layout="wide")

# Aspect Ratio Selection
aspect_ratios = {
    "16:9 (1280x720)": (1280, 720),
    "1:1 (1080x1080)": (1080, 1080),
    "9:16 (720x1280)": (720, 1280),
    "2:3 (800x1200)": (800, 1200),
}

# API key from environment variable
API_KEY = os.environ.get("NEBIUS_API_KEY")
if not API_KEY:
    st.error("API key not found. Please set the `NEBIUS_API_KEY` environment variable.")

# Function to call Nebius API for prompt generation
def generate_caption(image_base64, api_key):
    api_url = "https://api.studio.nebius.ai/v1/chat/completions"
    headers = {"Authorization": f"Bearer {api_key}"}
    payload = {
        "model": "mistralai/Mistral-Small-3.1-24B-Instruct-2503",
        "messages": [
            {"role": "system", "content": "You are an image-to-text prompt converter for AI image generation."},
            {"role": "user", "content": [
                {"type": "text", "text": "Describe This Image In Detail under 100 words in this format: [image content/subject, description of action, state, and mood], [art form, style], [artist/photographer reference if needed], [additional settings such as camera and lens settings, lighting, colors, effects, texture, background, rendering]."},
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
            ]},
        ],
        "temperature": 0.6,
    }

    try:
        response = requests.post(api_url, json=payload, headers=headers)
        if response.status_code == 200:
            return response.json()
        else:
            st.error(f"API Error: {response.status_code}, {response.text}")
            return {"error": response.text}
    except Exception as e:
        st.error(f"An exception occurred: {e}")
        return {"error": str(e)}

# File uploader for image
uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
selected_aspect = st.selectbox("Select Aspect Ratio", list(aspect_ratios.keys()))
width, height = aspect_ratios[selected_aspect]

if uploaded_image and API_KEY:
    image = Image.open(uploaded_image)
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    image_base64 = base64.b64encode(buffered.getvalue()).decode()
    
    st.image(image, caption="Uploaded Image", use_container_width=True)
    
    if st.button("Generate Image", use_container_width=True):
        st.write("Wait AI Is Generating Image...")
        result = generate_caption(image_base64, API_KEY)
        
        if "error" in result:
            st.error(f"Error: {result['error']}")
        else:
            try:
                caption = result.get("choices", [{}])[0].get("message", {}).get("content", "No caption generated.")
                
                # Construct the image generation URL
                image_url = f"https://image.pollinations.ai/prompt/{requests.utils.quote(caption)}?width={width}&height={height}&nologo=true&model=flux&enhance=true&seed=119"
                
                st.image(image_url, caption="Generated Image", use_container_width=True)
                
                # Download button for the generated image
                st.download_button(
    label="Download Image",
    data=requests.get(image_url).content,  # Fetch the image data
    file_name="generated_image.png",
    mime="image/png",
    use_container_width=True
)

            
            except Exception as e:
                st.error(f"Error processing the response: {e}")
else:
    st.info("Please upload an image and select an aspect ratio.")