import os
import io
import gradio as gr
from PIL import Image
from huggingface_hub import InferenceClient

# ---------------- GET HF TOKEN ----------------
HF_TOKEN = os.environ.get("HF_TOKEN")

if not HF_TOKEN:
    raise ValueError(
        "HF_TOKEN missing. Go to Hugging Face Space → Settings → Secrets → add HF_TOKEN."
    )

# Use official HF inference provider
client = InferenceClient(provider="hf-inference", api_key=HF_TOKEN)


# ---------------- STYLE PRESETS ----------------
STYLES = {
    "Realistic": "ultra realistic, 4k, cinematic lighting",
    "Anime": "anime style, vibrant colors",
    "Fantasy": "fantasy art, magical",
    "Cyberpunk": "cyberpunk neon futuristic",
}


# ---------------- GENERATE IMAGE ----------------
def generate_image(prompt, negative_prompt, style):

    if not prompt:
        raise gr.Error("Please enter a prompt.")

    # apply style text
    if style in STYLES:
        prompt = f"{prompt}, {STYLES[style]}"

    try:
        # FREE working model on CPU Spaces
        image_bytes = client.text_to_image(
            prompt=prompt,
            negative_prompt=negative_prompt,
            model="black-forest-labs/FLUX.1-schnell",
        )

        # convert → PIL
        if isinstance(image_bytes, bytes):
            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        else:
            image = image_bytes

        # save for download
        save_path = "/tmp/generated.png"
        image.save(save_path)

        return image, save_path

    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")


# ---------------- MODERN UI ----------------
custom_css = """
body {
    background: linear-gradient(135deg, #1e3c72, #2a5298);
}

.gradio-container {
    max-width: 900px !important;
    margin: auto;
    border-radius: 20px;
    padding: 25px;
    background: rgba(255,255,255,0.15);
    backdrop-filter: blur(18px);
}

h1, p, label { color: white !important; text-align: center; }
button { border-radius: 12px !important; font-weight: bold; }
"""


# ---------------- GRADIO APP ----------------
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:

    gr.Markdown("# 🎨 AI Image Generator")
    gr.Markdown("Fully working free Hugging Face image generator 🚀")

    prompt = gr.Textbox(label="Prompt", placeholder="A majestic lion in sunlight")

    negative_prompt = gr.Textbox(label="Negative Prompt (optional)")

    style = gr.Dropdown(list(STYLES.keys()), value="Realistic", label="Art Style")

    btn = gr.Button("✨ Generate Image")

    out_img = gr.Image(label="Generated Image")
    out_file = gr.File(label="Download Image")

    btn.click(generate_image, [prompt, negative_prompt, style], [out_img, out_file])


if __name__ == "__main__":
    demo.launch()