import gradio as gr
import requests
import os
import base64
import json
from io import BytesIO
from PIL import Image

# --- Configuration ---
# Your Hugging Face Space URL. Set this in your Space secrets for best practice.
YOUR_SITE_URL = os.environ.get("YOUR_SITE_URL", "https://huggingface.co/spaces/broadfield-dev/nano-banana")
YOUR_SITE_NAME = "Gemini 2.5 Image Editor Demo"
# Fallback API key from Hugging Face secrets
FALLBACK_API_KEY = os.environ.get("OPENROUTER_API_KEY")

def edit_image(image, prompt, api_key):
    """
    Sends an image and a prompt to the OpenRouter API and returns the edited image.
    """
    final_api_key = api_key if api_key else FALLBACK_API_KEY

    if not final_api_key:
        raise gr.Error("API Key is required. Please enter your OpenRouter API key or set it in the Space secrets.")

    if image is None:
        raise gr.Error("An input image is required. Please upload an image.")

    # Convert the input PIL image to a base64 string
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

    headers = {
        "Authorization": f"Bearer {final_api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": YOUR_SITE_URL,
        "X-Title": YOUR_SITE_NAME,
    }

    payload = {
        # *** CORRECTED MODEL NAME ***
        "model": "google/gemini-2.5-flash-image-preview:free",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/png;base64,{img_base64}"},
                    },
                ],
            }
        ],
        "max_tokens": 2048,
    }

    try:
        response = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=headers,
            data=json.dumps(payload)
        )
        response.raise_for_status()
        result_data = response.json()

        # --- *** CORRECTED RESPONSE PARSING LOGIC *** ---
        # The image is in the 'images' list within the message, not 'content'.
        message = result_data.get('choices', [{}])[0].get('message', {})
        
        if message and 'images' in message and message['images']:
            # Get the first image from the 'images' list
            image_data = message['images'][0]
            base64_string = image_data.get('image_url', {}).get('url', '')

            if base64_string and ',' in base64_string:
                # Remove the "data:image/png;base64," prefix
                base64_content = base64_string.split(',')[1]
                
                # Decode the base64 string and create a PIL image
                img_bytes = base64.b64decode(base64_content)
                edited_image = Image.open(BytesIO(img_bytes))
                return edited_image
            else:
                 raise gr.Error(f"API returned an invalid image format. Response: {json.dumps(result_data, indent=2)}")
        else:
            raise gr.Error(f"API did not return an image. Full Response: {json.dumps(result_data, indent=2)}")

    except requests.exceptions.HTTPError as err:
        error_body = err.response.text
        if err.response.status_code == 401:
            raise gr.Error("Authentication failed. Check your OpenRouter API key.")
        elif err.response.status_code == 429:
             raise gr.Error("Rate limit exceeded or insufficient credits. Check your OpenRouter account.")
        else:
            raise gr.Error(f"An API error occurred: {error_body}")
    except Exception as e:
        # Catch any other unexpected errors, like JSON parsing failures
        raise gr.Error(f"An unexpected error occurred: {str(e)}")

# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as iface:
    gr.Markdown(
        """
        # Image Editing with Google Gemini 2.5 Flash via OpenRouter
        OpenRouter API key included, but may run out of free inference, so you can use your own key.
        Try the chatbot version here (BYOK)
        [https://huggingface.co/spaces/broadfield-dev/nano-banana-chat](https://huggingface.co/spaces/broadfield-dev/nano-banana-chat)
        Upload an image, describe the edit you want.

        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Upload Image")
            prompt_input = gr.Textbox(label="Prompt", placeholder="e.g., 'make the background black and white' or 'add sunglasses to the person'")
            api_key_input = gr.Textbox(
                label="OpenRouter API Key (optional)",
                placeholder="Enter your key here (uses Space secret as fallback)",
                type="password"
            )
            submit_button = gr.Button("Generate Image", variant="primary")
        with gr.Column(scale=1):
            image_output = gr.Image(label="Edited Image")

    submit_button.click(
        fn=edit_image,
        inputs=[image_input, prompt_input, api_key_input],
        outputs=image_output
    )

iface.launch(ssr_mode=False)