import gradio as gr
import requests
import random
from io import BytesIO
from PIL import Image, ImageDraw
import os
import time

# Base URL from HF secret or fallback
BASE_URL = os.environ.get('POLLINATIONS_URL')  # e.g. "https://image.pollinations.ai/prompt/"

# Load watermark logo (CodeHubb logo)
WATERMARK_PATH = "CodeHubb.png"
watermark = Image.open(WATERMARK_PATH).convert("RGBA")
watermark = watermark.resize((150, 40))  # Resize as needed

def add_watermark(image):
    image = image.convert("RGBA")
    # Create a transparent layer and paste watermark
    layer = Image.new("RGBA", image.size, (0, 0, 0, 0))
    layer.paste(watermark, (10, 10))  # Top-left corner with 10px padding
    final_image = Image.alpha_composite(image, layer)
    return final_image.convert("RGB")

def generate_image(prompt, model):
    if not prompt:
        raise gr.Error("Please enter a prompt.")

    seed = random.randint(1, 999999)
    url = f"{BASE_URL}{prompt}?width=2048&height=2048&seed={seed}&nologo=true&model={model}"

    max_retries = 2
    for attempt in range(max_retries):
        try:
            response = requests.get(url, stream=True)
            response.raise_for_status()
            img = Image.open(BytesIO(response.content))
            watermarked_img = add_watermark(img)
            return watermarked_img

        except requests.exceptions.HTTPError as e:
            if response.status_code == 500 and 'Access to kontext model' in response.text:
                if attempt < max_retries - 1:
                    time.sleep(1)
                    continue
                raise gr.Error("Access denied for kontext model (limit reached). Try turbo/flux or authenticate at pollinations.ai.")
            else:
                raise gr.Error(f"Error: {response.status_code} - {response.text}")
        except Exception as e:
            raise gr.Error(f"Unexpected error: {str(e)}")

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# CodeHubb AI Image Generator\n**Design. Develop. Dominate.**")

    prompt_input = gr.Textbox(label="Prompt", placeholder="e.g., A futuristic cityscape at night")
    model_input = gr.Dropdown(choices=["kontext", "turbo", "flux"], label="Model", value="turbo")

    generate_btn = gr.Button("Generate")
    output_image = gr.Image(label="Generated Image")

    generate_btn.click(generate_image, inputs=[prompt_input, model_input], outputs=output_image)

demo.queue()
demo.launch()