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

API_KEY = os.getenv("NVIDIA_API_KEY")
invoke_url = "https://ai.api.nvidia.com/v1/genai/stabilityai/sdxl-turbo"
headers = {
    "Authorization": f"Bearer {API_KEY}",
    "Accept": "application/json",
}

def generate_kindle_cover(prompt):
    print(f"Generating cover for prompt: {prompt}")
    print(f"API Key (first 5 characters): {API_KEY[:5]}..." if API_KEY else "API Key is not set")

    payload = {
        "text_prompts": [{"text": prompt}],
        "seed": 0,
        "sampler": "K_EULER_ANCESTRAL",
        "steps": 2
    }

    try:
        print("Sending request to NVIDIA API...")
        response = requests.post(invoke_url, headers=headers, json=payload)
        print(f"Received response with status code: {response.status_code}")
        print(f"Response headers: {json.dumps(dict(response.headers), indent=2)}")

        if response.status_code == 401:
            return "Error: Unauthorized access (401). Please check your API key.", None
        elif response.status_code != 200:
            return f"Error: Received status code {response.status_code} from API. Response: {response.text}", None

        try:
            response_body = response.json()
            print("Response Body:", json.dumps(response_body, indent=2))
            
            # Extract and decode the image data
            image_data = base64.b64decode(response_body['results'][0]['image'])
            image = Image.open(io.BytesIO(image_data))
            
            return "Success: Image generated successfully.", image

        except json.JSONDecodeError:
            return f"Error: Failed to parse JSON response. Raw response: {response.text}", None
        except KeyError:
            return "Error: Expected image data not found in the response.", None

    except Exception as e:
        print(f"An error occurred: {str(e)}")
        return f"Error: {str(e)}", None

iface = gr.Interface(
    fn=generate_kindle_cover,
    inputs="text",
    outputs=["text", gr.Image()],
    title="Kindle Cover Generator",
    description="Generate high-quality covers for Amazon Kindle books using the NVIDIA SDXL-Turbo model."
)

iface.launch(debug=True)