from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import os
import subprocess
from huggingface_hub import InferenceClient
from io import BytesIO
from PIL import Image
import re

# Initialize the Flask app
app = Flask(__name__)
CORS(app)  # Enable CORS for all routes

# Initialize the InferenceClient with your Hugging Face token
HF_TOKEN = os.environ.get("HF_TOKEN")  # Ensure to set your Hugging Face token in the environment
client = InferenceClient(token=HF_TOKEN)

# Hardcoded negative prompt
NEGATIVE_PROMPT_FINGERS = """2D,missing fingers, extra fingers, elongated fingers, fused fingers, 
mutated fingers, poorly drawn fingers, disfigured fingers, 
too many fingers, deformed hands, extra hands, malformed hands, 
blurry hands, disproportionate fingers,sexual,nudity,intimacy view"""

# Define the list of explicit keywords

EXPLICIT_KEYWORDS = [
    # General sexual content
    "sexual", "sex", "sexualized", "sexually", "sexist","sexy","sexily","sexious","sexioud","boobs", "boob", "breasts","breast", "cleavage","cleavages","nipples","nipple","tits","tit","penis","penises","penes","phallus", "porn", "pornography", "adult video", "hentai", "ecchi", "fetish", "nude", "nudes","nudity","vagina","provocative", "obscene", "vulgar", "intimate", "kinky", "hardcore", "softcore","seduce", 
    "seductive", "sensual", "explicit", "explicitly","lewd", "taboo", "NSFW", "erotic", "erotica","erotical","erotically","arousal","arouse","lust","hot","hottest","hotties","hottier","hotily","hotious","pornstar", "adult actor", "adult actors", "adult actress", "adult actresses","silhouette","silhouettes","inappropriate","inappropriates","condom","condoms","condomes",

    # Sexual acts and positions
    "threesome", "orgy", "masturbation", "masturbate", "penetration", "intercourse", 
    "oral sex", "blowjob", "handjob", "anal sex", "doggy style", "69", "orgasm", "cum", 
    "ejaculate", "ejaculation", "sperm", "semen", "cumming", "squirting", 

    # Body parts
    "genital", "genitals", "vagina", "vaginal", "anus", "anal", "butt", "buttocks", "butthole", 
    "ass", "prostate", "erection", "clitoris", "pussy", "dick", "balls", "testicles","touching body part", "touching body","showing body","touching body parts","showing body part","showing body parts","woman body", "man body", 

    # Clothing and behavior
    "naked", "bare", "lingerie", "thong", "striptease", "stripper", "bikini", "topless", 
    "undress", "undressing", "revealing", "skimpy", "suggestive", "transparent clothing", 
    "provocative pose","sheer","exposed","innerwear", "innerwears","innerweares","undies", "undy", "underwear", "underwears", 
"underweares", "panty", "panties", "brief", "briefs", "hipster", "hipsters", "hipsteres", "hollow cut-out", "hollow cut-outs", "hollow cut-outes", "cut-out", "cut-outs", "cut-outes","bra","bras",

    # BDSM and fetishes
    "bdsm", "dominatrix", "submission", "bondage", "whip", "chains", "gag", "spanking", 
    "fetish", "roleplay", "taboo", "voyeur", "exhibitionist", "peeping", "discipline", 
    "dominant", "submissive", "consensual non-consent",

    # Sex toys and related
    "dildo", "sex toy", "vibrator", "butt plug", "strap-on", "lube", "latex clothing",

    # Derogatory and explicit language
    "fuck", "fucking", "fucker", "fuckers", "slut", "whore", "prostitute", "hooker", 
    "escort", "camgirl", "camwhore", "sugar daddy", "sugar baby", "cumslut", "milf", 
    "daddy kink", "mommy kink", 

    # Sexual violence and illegal content
    "abuse", "violence", "rape", "sexual violence", "molestation", "pedophilia", 
    "child porn", "underage", "illegal content", "incest", 

    # Mental health and harmful themes
    "suicide", "self-harm", "depression", "kill myself", "worthless", 
    "abuse victim", "emotional abuse", "manipulation", "gore", "violent content"
]

# Function to scan the entire prompt for explicit keywords
def scan_prompt(prompt, keywords):
    # Create a regex pattern to match any keyword (case insensitive)
    pattern = r'\b(?:' + '|'.join(re.escape(keyword) for keyword in keywords) + r')\b'
    # Find all matches in the prompt
    matches = re.findall(pattern, prompt, flags=re.IGNORECASE)
    # Return True if matches are found, and the list of matched keywords
    return bool(matches), matches

@app.route('/')
def home():
    return "Welcome to the Image Background Remover!"

# Function to generate an image from a text prompt
def generate_image(prompt, negative_prompt=None, height=512, width=512, model="stabilityai/stable-diffusion-2-1", num_inference_steps=50, guidance_scale=7.5, seed=None):
    try:
        # Generate the image using Hugging Face's inference API with additional parameters
        image = client.text_to_image(
            prompt=prompt, 
            negative_prompt=NEGATIVE_PROMPT_FINGERS,
            height=height, 
            width=width, 
            model=model,
            num_inference_steps=num_inference_steps,  # Control the number of inference steps
            guidance_scale=guidance_scale,  # Control the guidance scale
            seed=seed  # Control the seed for reproducibility
        )
        return image  # Return the generated image
    except Exception as e:
        print(f"Error generating image: {str(e)}")
        return None

# Flask route for the API endpoint to generate an image
@app.route('/generate_image', methods=['POST'])
def generate_api():
    data = request.get_json()

    # Extract required fields from the request
    prompt = data.get('prompt', '')
    negative_prompt = data.get('negative_prompt', None)
    height = data.get('height', 1024)  # Default height
    width = data.get('width', 720)  # Default width
    num_inference_steps = data.get('num_inference_steps', 50)  # Default number of inference steps
    guidance_scale = data.get('guidance_scale', 7.5)  # Default guidance scale
    model_name = data.get('model', 'stabilityai/stable-diffusion-2-1')  # Default model
    seed = data.get('seed', None)  # Seed for reproducibility, default is None

    if not prompt:
        return jsonify({"error": "Prompt is required"}), 400

    try:
        # Check for explicit content using scan_prompt
        is_nsfw, found_keywords = scan_prompt(prompt, EXPLICIT_KEYWORDS)
        if is_nsfw:
            print(f"Explicit keywords found: {found_keywords}")
            # Return the pre-defined "nsfw.jpg" image
            return send_file(
                "nsfw.jpg",
                mimetype='image/png',
                as_attachment=False,
                download_name='nsfw_detected.png'
            )

        # Call the generate_image function with the provided parameters
        image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)

        if image:
            # Save the image to a BytesIO object
            img_byte_arr = BytesIO()
            image.save(img_byte_arr, format='PNG')  # Convert the image to PNG
            img_byte_arr.seek(0)  # Move to the start of the byte stream

            # Send the generated image as a response
            return send_file(
                img_byte_arr, 
                mimetype='image/png', 
                as_attachment=False,  # Send the file as an attachment
                download_name='generated_image.png'  # The file name for download
            )
        else:
            return jsonify({"error": "Failed to generate image"}), 500
    except Exception as e:
        print(f"Error in generate_api: {str(e)}")  # Log the error
        return jsonify({"error": str(e)}), 500

# Add this block to make sure your app runs when called
if __name__ == "__main__":
    subprocess.Popen(["python", "wk.py"])  # Start awake.py

    app.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing