import gradio as gr
import requests
from PIL import Image
import io

def generate_kontext_image(input_image, prompt, width=1024, height=1024, seed=-1, model="dreamshaper", nologo=True, enhance=False):
    """
    Generate a transformed image using the Pollinations API.
    
    Args:
        input_image (PIL.Image): Input image to transform.
        prompt (str): Prompt for the transformation.
        width (int): Width of the output image.
        height (int): Height of the output image.
        seed (int): Random seed for generation (-1 for random).
        model (str): Model to use (default: 'dreamshaper').
        nologo (bool): Whether to exclude logo.
        enhance (bool): Whether to enhance the image.
    
    Returns:
        PIL.Image or str: Generated image or error message.
    """
    # Step 1: Convert the input PIL Image to bytes for uploading.
    image_bytes = io.BytesIO()
    input_image.save(image_bytes, format='JPEG')
    image_bytes.seek(0)
    
    input_image_url = ""
    # Step 2: Upload the image bytes to get a public URL.
    try:
        upload_response = requests.post(
            'https://image.pollinations.ai/upload',
            files={'file': ('input_image.jpg', image_bytes, 'image/jpeg')}
        )
        upload_response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
        upload_result = upload_response.json()
        input_image_url = upload_result.get('ipfs')
        
        if not input_image_url:
            return "Error: Could not retrieve a public URL after uploading the image."
            
    except requests.RequestException as e:
        return f"Error: Failed to upload the image to the server - {e}"

    # Step 3: Use the public URL to make the final image generation request.
    base_url = "https://image.pollinations.ai/prompt"
    
    # URL-encode the prompt to handle spaces and special characters
    encoded_prompt = requests.utils.quote(prompt)
    api_url = f"{base_url}/{encoded_prompt}"

    query_params = {
        "model": model,
        "image": input_image_url,
        "width": width,
        "height": height,
        "seed": seed,
        "nologo": str(nologo).lower(),
        "enhance": str(enhance).lower()
    }

    try:
        # Make the API request
        response = requests.get(api_url, params=query_params, stream=True)
        response.raise_for_status()
        
        # Convert response content to a PIL Image and return it
        output_image = Image.open(io.BytesIO(response.content))
        return output_image
        
    except requests.RequestException as e:
        error_details = str(e)
        try:
            # Try to get a more specific error message from the API response
            error_details = e.response.json().get("message", e.response.text)
        except:
            pass # Keep the original exception text if parsing fails
        return f"Error: API request failed. Details: {error_details}"


def app_interface(input_image, prompt, width, height, seed, nologo, enhance):
    """
    Gradio interface function to handle user inputs and display results.
    """
    if input_image is None:
        return "Please upload an image."
    if not prompt:
        return "Please provide a prompt."
    
    # Using 'dreamshaper' as it's a free and effective model for this task
    return generate_kontext_image(
        input_image=input_image,
        prompt=prompt,
        width=width,
        height=height,
        seed=seed,
        model="dreamshaper",
        nologo=nologo,
        enhance=enhance
    )

# Define the Gradio interface
with gr.Blocks(title="Image Transformation") as demo:
    gr.Markdown("# Image Transformation App")
    gr.Markdown("Upload an image, provide a transformation prompt, and generate a new image.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload Image")
            prompt = gr.Textbox(label="Prompt", placeholder="e.g., transform this image into a surreal painting")
            width = gr.Slider(minimum=256, maximum=2048, value=1024, step=1, label="Width")
            height = gr.Slider(minimum=256, maximum=2048, value=1024, step=1, label="Height")
            seed = gr.Number(value=-1, label="Seed (-1 for random)", precision=0)
            nologo = gr.Checkbox(value=True, label="No Logo")
            enhance = gr.Checkbox(value=False, label="Enhance Image")
            submit_button = gr.Button("Generate Image")
        
        with gr.Column():
            output = gr.Image(label="Generated Image")
    
    submit_button.click(
        fn=app_interface,
        inputs=[input_image, prompt, width, height, seed, nologo, enhance],
        outputs=output
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()