import gradio as gr
import requests
import os
import io
from PIL import Image
from typing import Optional


# Load Hugging Face Token from environment variables
def load_environment():
    return os.getenv("HF_TOKEN")


# Enhance the text prompt for better photorealistic results
def craft_realistic_prompt(base_prompt: str) -> str:
    realistic_modifiers = [
        "photorealistic", "high resolution", "sharp focus",
        "professional photography", "natural lighting", "detailed textures"
    ]
    return f"{' '.join(realistic_modifiers)}, {base_prompt}, shot on professional camera, 8k resolution"


# Query Hugging Face API to generate image from text
def query_hf_api(prompt: str) -> Optional[bytes]:
    model_url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
    HF_TOKEN = load_environment()

    if not HF_TOKEN:
        raise ValueError("Hugging Face token not found. Set HF_TOKEN in environment variables.")

    headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}
    payload = {
        "inputs": craft_realistic_prompt(prompt),
        "parameters": {
            "negative_prompt": "cartoon, anime, low quality, bad anatomy, blurry, unrealistic, painting, drawing, sketch",
            "num_inference_steps": 75,
            "guidance_scale": 8.5,
        }
    }

    response = requests.post(model_url, headers=headers, json=payload, timeout=120)
    response.raise_for_status()
    return response.content


# Generate image from text and return as PIL image
def generate_image(prompt: str):
    try:
        image_bytes = query_hf_api(prompt)
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        return image
    except Exception as e:
        return f"Error: {str(e)}"


# Modern Gradio UI using Blocks
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 Wangoes Text-to-Image Generator")
    gr.Markdown("Generate stunning AI images from text using Hugging Face models.")

    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Enter your prompt",
                                      placeholder="Describe the image you want to generate...")
            generate_button = gr.Button("Generate Image", variant="primary")

        with gr.Column():
            image_output = gr.Image(label="Generated Image", type="pil")

    gr.Examples(
        examples=[
            "A futuristic cityscape at sunset with flying cars",
            "A serene mountain lake with crystal clear water",
            "A cozy cabin in the woods during winter"
        ],
        inputs=prompt_input
    )

    generate_button.click(fn=generate_image, inputs=prompt_input, outputs=image_output)

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