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="Image Caption Generator", layout="wide")

# 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 resize the image
def resize_image(image, max_width=300):
    """
    Resize the image to a fixed width (300) while maintaining the aspect ratio.
    """
    width_percent = max_width / float(image.size[0])
    new_height = int((float(image.size[1]) * float(width_percent)))
    resized_image = image.resize((max_width, new_height), Image.Resampling.LANCZOS)
    return resized_image

# Function to compress the image before converting to Base64
def compress_image(image, quality=70):
    """
    Compress the image to reduce the size before Base64 encoding.
    """
    buffered = BytesIO()
    image.save(buffered, format="JPEG", quality=quality)  # Save as JPEG with compression
    return buffered.getvalue()

# Function to convert image to Base64
def convert_to_base64(image_data):
    """
    Convert the compressed image to Base64 encoding.
    """
    image_base64 = base64.b64encode(image_data).decode()
    return image_base64

# Function to call Nebius API
def generate_caption(image_base64, api_key):
    """
    Calls the Nebius API with the Base64-encoded image and generates a caption.
    """
    api_url = "https://api.studio.nebius.ai/v1/chat/completions"
    headers = {"Authorization": f"Bearer {api_key}"}
    payload = {
        "model": "llava-hf/llava-1.5-13b-hf",
        "messages": [
            {
                "role": "system",
                "content": """observe each and every detail of the image and craft a detailed prompt under 75 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].""",
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Describe This Image In Detail"},
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
                ],
            },
        ],
        "temperature": 0,
    }

    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)}



uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_image:
    # Load the uploaded image
    original_image = Image.open(uploaded_image)

    # Resize the image to a fixed width of 300px
    resized_image = resize_image(original_image)

    # Display the resized image
    st.image(resized_image, caption="Resized Image", use_container_width=True)

    # Compress the image before converting to Base64
    compressed_image_data = compress_image(resized_image)

    # Convert compressed image to Base64
    base64_string = convert_to_base64(compressed_image_data)

    # Call the Nebius API to generate a caption
    if st.button("Generate Prompt", use_container_width=True):
        st.write("Generating Prompt...")
        result = generate_caption(base64_string, API_KEY)

        # Display the generated caption
        if "error" in result:
            st.error(f"Error: {result['error']}")
        else:
            try:
                caption = (
                    result.get("choices", [{}])[0]
                    .get("message", {})
                    .get("content", "No caption generated.")
                )
                
                st.text_area("Prompt Generated", caption, height=200, key="caption_output")
            except Exception as e:
                st.error(f"Error processing the response: {e}")
else:
    if not API_KEY:
        st.warning("Please set the `NEBIUS_API_KEY` environment variable.")
    else:
        st.info("Please upload an image.")

# Style adjustments
st.markdown(
    """
    <style>
     
        .css-1d391kg {text-align: center;}
        .css-1v3fvcr {justify-content: center; display: flex;}
        .css-1y4ccs5 {margin: 0 auto;}
        .css-1aumxhk {display: flex; justify-content: center;}
        .css-1k6kn8p {justify-content: center; align-items: center; display: flex;}
        .css-ffhzg6 {text-align: center;}
        textarea {
            color: white !important;
            background-color: #262626 !important;
        }
    
        textarea {
            color: white !important;
            background-color: #262626 !important;
        }
    </style>
    """, unsafe_allow_html=True
)