import gradio as gr
import requests
import random
import os

def process(prompt):
    """
    Processes the user's text prompt to either generate an image or display a content warning.
    """
    # Expanded list of restricted words.
    restricted_words = [
        "sex", "erotic", "nude", "explicit", "porn", "xxx", "18+", "nudity", "sexual",
        "lewd", "obscene", "hentai", "bondage", "bdsm", "fetish", "voyeur",
        "masturbate", "orgasm", "gangbang", "rape", "incest", "pedophile", "child abuse",
        "bestiality", "prostitution", "threesome", "anal", "oral", "intercourse",
        "ejaculation", "cumshot", "hardcore", "softcore", "lingerie", "stripper",
        "dildo", "vibrator", "panty", "bra", "nipple", "genitals", "pubic", "ass", "tits",
        "dick", "pussy", "blowjob", "handjob", "deepthroat", "g-string", "bikini",
        "topless", "bottomless", "kinky", "dominant", "submissive", "sadism", "masochism",
        "orgies", "sexually suggestive", "vulgar"
    ]

    # Check if the user's prompt contains any of the restricted words (case-insensitive).
    if any(word in prompt.lower() for word in restricted_words):
        return "This is mature content. It's not permitted to generate images on this topic."

    # If the prompt is safe, proceed with image generation.
    filename = str(random.randint(111111111, 999999999)) + ".png"
    file_path = os.path.join(os.path.dirname(__file__), filename)
    
    # Use Pollinations API to get the image.
    response = requests.get(
        "https://image.pollinations.ai/prompt/" + prompt + 
        "?model=flux-realism&width=2048&height=2048&nologo=true&seed=" + 
        str(random.randint(0, 999999999))
    )

    if response.status_code == 200:
        with open(file_path, "wb") as f:
            f.write(response.content)
        return file_path
    else:
        # Return an error message if the API call fails.
        return f"Error: Could not retrieve image. Status code: {response.status_code}"

# Define the Gradio interface.
title = "Pollinations Image Generator"
description = "This app generates images from text using the Pollinations API."
article = "Note: This app is filtered to prevent the generation of mature content."

iface = gr.Interface(
    fn=process, 
    inputs=gr.Textbox(
        lines=2, 
        placeholder="Enter your prompt here...", 
        label="Prompt"
    ), 
    outputs=gr.Image(
        type="filepath", 
        label="Generated Image"
    ), 
    title=title, 
    description=description,
    article=article
)

iface.launch(server_name="0.0.0.0", server_port=7860, pwa=True)