import os
import gradio as gr
from PIL import Image
import io
import tempfile
import torch
from diffusers import StableDiffusionXLPipeline

# ============================================
# SEO Configuration for Hugging Face Spaces
# ============================================
TITLE = "FreeAI PicGen - Free Unlimited AI Image Generator"
DESCRIPTION = """
🎨 **FreeAI PicGen** - Generate stunning AI images for FREE using Stable Diffusion XL!

### ✨ Features
- 📝 **Text-to-Image**: Create images from text descriptions (512x512 default)
- 💡 **Smart Prompt Helper**: Auto-optimize your prompts for better results
- ⚡ **Unlimited Free Generations**: No limits, powered by Hugging Face Inference API
- 💾 **Instant Download**: One-click download for all generated images

### 🔒 Privacy First
Your images are processed securely and not stored permanently.

🚀 **Powered by**: [Hugging Face Diffusers](https://huggingface.co/docs/diffusers) | [Stable Diffusion XL](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
"""

# ============================================
# Configuration
# ============================================
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    print("⚠️ Warning: HF_TOKEN environment variable not set. Please set it in your Space secrets.")

# Load the model pipeline
MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
pipe = StableDiffusionXLPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    use_safetensors=True,
    token=HF_TOKEN
)
# pipe.enable_model_cpu_offload()  # Save VRAM

# Prompt optimization templates
PROMPT_TEMPLATES = {
    "quality": "masterpiece, best quality, highly detailed, {prompt}, 8k uhd, dslr, soft lighting, film grain, sharp focus",
    "cinematic": "cinematic shot, {prompt}, depth of field, motion blur, cinematic lighting, film grain, 35mm, unreal engine 5, octane render",
    "anime": "anime style, {prompt}, studio ghibli, detailed face, vibrant colors, anime screencap, cel shaded",
    "digital_art": "digital art, {prompt}, concept art, sharp focus, illustration, trending on artstation, by artgerm and greg rutkowski",
    "photorealistic": "professional photography, {prompt}, photorealistic, hyperrealistic, detailed skin texture, professional lighting, RAW photo, 8k uhd"
}

# ============================================
# Core Functions
# ============================================

def optimize_prompt(prompt, style="quality"):
    """Optimize user prompt using templates"""
    if not prompt or not prompt.strip():
        return ""
    template = PROMPT_TEMPLATES.get(style, PROMPT_TEMPLATES["quality"])
    return template.format(prompt=prompt.strip())

def save_image_to_temp(image, prefix="generated"):
    """Save PIL image to temporary file and return path"""
    if image is None:
        return None
    temp_dir = tempfile.gettempdir()
    temp_path = os.path.join(temp_dir, f"{prefix}_{hash(str(image))}.png")
    image.save(temp_path, format='PNG')
    return temp_path

def generate_text2image(prompt, negative_prompt, width, height, steps, cfg_scale, seed):
    """Generate image from text using local SDXL pipeline"""
    try:
        if not prompt or not prompt.strip():
            return None, "❌ Error: Please enter a prompt!", None
        
        # Generate image
        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
        if seed != -1:
            generator.manual_seed(int(seed))
        
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=int(width),
            height=int(height),
            num_inference_steps=int(steps),
            guidance_scale=float(cfg_scale),
            generator=generator
        ).images[0]
        
        # Save to temp file for download
        download_path = save_image_to_temp(image, "txt2img")
        
        return image, "✅ Image generated successfully!", download_path
        
    except Exception as e:
        error_str = str(e)
        if "out of memory" in error_str.lower():
            return None, "⏳ GPU out of memory. Try lower resolution or fewer steps.", None
        elif "cuda" in error_str.lower():
            return None, "⏳ CUDA error. Ensure GPU is available.", None
        else:
            return None, f"❌ Error: {error_str}", None

# def generate_img2img(prompt, input_image, strength, negative_prompt, steps, cfg_scale, seed):
#     """Generate image from image using HF Inference API"""
#     try:
#         if not prompt or not prompt.strip():
#             return None, "❌ Error: Please enter a prompt!", None
        
#         if input_image is None:
#             return None, "❌ Error: Please upload an image!", None
        
#         # Try to generate image from image
#         try:
#             image = client.image_to_image(
#                 image=input_image,
#                 prompt=prompt,
#                 negative_prompt=negative_prompt,
#                 model=MODEL_ID,
#                 strength=float(strength),
#                 num_inference_steps=int(steps),
#                 guidance_scale=float(cfg_scale),
#                 seed=None if seed == -1 else int(seed)
#             )
            
#             # Save to temp file for download
#             download_path = save_image_to_temp(image, "img2img")
            
#             return image, "✅ Image transformed successfully!", download_path
            
#         except Exception as api_error:
#             error_msg = str(api_error)
#             if "not supported" in error_msg.lower() or "task" in error_msg.lower():
#                 return None, "⚠️ Image-to-Image requires Hugging Face Pro Inference API. Text-to-Image is available for free!", None
#             raise api_error
            
#     except Exception as e:
#         error_str = str(e)
#         if "rate limit" in error_str.lower():
#             return None, "⏳ Rate limit reached. Please wait 30 seconds and try again.", None
#         elif "401" in error_str or "unauthorized" in error_str.lower():
#             return None, "🔑 Authentication error. Please set HF_TOKEN environment variable.", None
#         elif "503" in error_str or "loading" in error_str.lower():
#             return None, "⏳ Model is loading. Please wait a moment and retry.", None
#         else:
#             return None, f"❌ Error: {error_str}", None

# ============================================
# Custom CSS with AdSense Placeholder
# ============================================
CUSTOM_CSS = """
/* =========================================
   GOOGLE ADSENSE PLACEHOLDER
   Uncomment below and add your publisher ID
   ========================================= */
/*
<script async src="https://pagead2.googlesyndication.com/pagead/js/adsbygoogle.js?client=ca-pub-YOUR_PUBLISHER_ID" crossorigin="anonymous"></script>
<!-- Ad Unit -->
<ins class="adsbygoogle"
     style="display:block"
     data-ad-client="ca-pub-YOUR_PUBLISHER_ID"
     data-ad-slot="YOUR_AD_SLOT_ID"
     data-ad-format="auto"
     data-full-width-responsive="true"></ins>
<script>
     (adsbygoogle = window.adsbygoogle || []).push({});
</script>
*/

.gradio-container {
    font-family: 'Inter', system-ui, -apple-system, sans-serif;
}

.title-container {
    text-align: center;
    margin-bottom: 2rem;
}

.title-container h1 {
    color: #6366f1;
    font-size: 2.5rem;
    margin-bottom: 0.5rem;
}

.title-container p {
    color: #6b7280;
    font-size: 1.1rem;
}

.generate-btn {
    background: linear-gradient(135deg, #6366f1 0%, #8b5cf6 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
}

.generate-btn:hover {
    transform: translateY(-1px);
    box-shadow: 0 4px 12px rgba(99, 102, 241, 0.4);
}

.pro-badge {
    background: linear-gradient(135deg, #f59e0b 0%, #d97706 100%);
    color: white;
    padding: 2px 8px;
    border-radius: 4px;
    font-size: 0.75rem;
    font-weight: 600;
    margin-left: 8px;
}
"""

# ============================================
# Gradio Interface
# ============================================

with gr.Blocks() as demo:
    
    # Header
    with gr.Column(elem_classes="title-container"):
        gr.Markdown("# 🎨 FreeAI PicGen")
        gr.Markdown(DESCRIPTION)
    
    # Google AdSense Banner
    gr.HTML("""
    <!-- Google AdSense Placeholder -->
    <div style="text-align: center; margin: 20px 0;">
      <script async src="https://pagead2.googlesyndication.com/pagead/js/adsbygoogle.js?client=ca-pub-YOUR_PUBLISHER_ID" crossorigin="anonymous"></script>
      <ins class="adsbygoogle"
           style="display:block"
           data-ad-client="ca-pub-YOUR_PUBLISHER_ID"
           data-ad-slot="YOUR_AD_SLOT_ID"
           data-ad-format="auto"
           data-full-width-responsive="true"></ins>
      <script>
           (adsbygoogle = window.adsbygoogle || []).push({});
      </script>
    </div>
    """)
    
    # Main Tabs
    with gr.Tabs():
        
        # ========================================
        # Tab 1: Text-to-Image
        # ========================================
        with gr.Tab("📝 Text to Image"):
            with gr.Row():
                # Left Column - Inputs
                with gr.Column(scale=1):
                    t2i_prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="A serene lake surrounded by mountains at sunset...",
                        lines=3,
                        max_lines=5
                    )
                    
                    t2i_negative = gr.Textbox(
                        label="Negative Prompt (Optional)",
                        placeholder="blur, low quality, distorted, watermark, signature...",
                        lines=2
                    )
                    
                    # Prompt Helper Section
                    with gr.Row():
                        style_dropdown = gr.Dropdown(
                            choices=list(PROMPT_TEMPLATES.keys()),
                            value="quality",
                            label="Style Preset",
                            info="Select a style to optimize your prompt"
                        )
                        optimize_btn = gr.Button("✨ Optimize Prompt", variant="secondary", size="sm")
                    
                    # Advanced Settings
                    with gr.Accordion("⚙️ Advanced Settings", open=False):
                        with gr.Row():
                            t2i_width = gr.Slider(
                                minimum=256, maximum=1024, value=512, step=64,
                                label="Width"
                            )
                            t2i_height = gr.Slider(
                                minimum=256, maximum=1024, value=512, step=64,
                                label="Height"
                            )
                        
                        with gr.Row():
                            t2i_steps = gr.Slider(
                                minimum=10, maximum=50, value=30, step=1,
                                label="Inference Steps"
                            )
                            t2i_cfg = gr.Slider(
                                minimum=1.0, maximum=15.0, value=7.5, step=0.5,
                                label="CFG Scale"
                            )
                        
                        t2i_seed = gr.Number(
                            value=-1,
                            label="Seed (-1 for random)",
                            precision=0
                        )
                    
                    t2i_generate = gr.Button(
                        "🎨 Generate Image",
                        variant="primary",
                        elem_classes="generate-btn",
                        size="lg"
                    )
                
                # Right Column - Outputs
                with gr.Column(scale=1):
                    t2i_output = gr.Image(
                        label="Generated Image",
                        type="pil"
                    )
                    t2i_status = gr.Textbox(
                        label="Status",
                        interactive=False
                    )
                    t2i_download = gr.File(
                        label="Download Image (PNG)",
                        file_types=[".png"]
                    )
            
            # Event Handlers
            optimize_btn.click(
                fn=optimize_prompt,
                inputs=[t2i_prompt, style_dropdown],
                outputs=t2i_prompt
            )
            
            t2i_generate.click(
                fn=generate_text2image,
                inputs=[t2i_prompt, t2i_negative, t2i_width, t2i_height, t2i_steps, t2i_cfg, t2i_seed],
                outputs=[t2i_output, t2i_status, t2i_download]
            )
        
        # ========================================
        # Tab 2: Image-to-Image (Pro Feature)
        # ========================================
        # with gr.Tab("🔄 Image to Image"):
        #     gr.Markdown("""
        #     ### 🔄 Image-to-Image Transformation
        #     <span class="pro-badge">PRO FEATURE</span>
            
        #     Transform existing images using AI. This feature requires a Hugging Face Pro Inference API subscription.
        #     Free tier users can use Text-to-Image which is unlimited!
        #     """)
            
        #     with gr.Row():
        #         # Left Column
        #         with gr.Column(scale=1):
        #             i2i_prompt = gr.Textbox(
        #                 label="Prompt",
        #                 placeholder="Describe what you want to transform the image into...",
        #                 lines=3
        #             )
                    
        #             i2i_negative = gr.Textbox(
        #                 label="Negative Prompt (Optional)",
        #                 placeholder="blur, low quality, distorted...",
        #                 lines=2
        #             )
                    
        #             i2i_input = gr.Image(
        #                 label="Upload Source Image",
        #                 type="pil",
        #                 sources=["upload", "clipboard"]
        #             )
                    
        #             i2i_strength = gr.Slider(
        #                 minimum=0.0, maximum=1.0, value=0.75, step=0.05,
        #                 label="Denoising Strength",
        #                 info="0.0 = original image, 1.0 = completely new image"
        #             )
                    
        #             with gr.Accordion("⚙️ Advanced Settings", open=False):
        #                 with gr.Row():
        #                     i2i_steps = gr.Slider(
        #                         minimum=10, maximum=50, value=30, step=1,
        #                         label="Inference Steps"
        #                     )
        #                     i2i_cfg = gr.Slider(
        #                         minimum=1.0, maximum=15.0, value=7.5, step=0.5,
        #                         label="CFG Scale"
        #                     )
        #                 i2i_seed = gr.Number(
        #                     value=-1,
        #                     label="Seed (-1 for random)",
        #                     precision=0
        #                 )
                    
        #             i2i_generate = gr.Button(
        #                 "🔄 Transform Image",
        #                 variant="primary",
        #                 elem_classes="generate-btn",
        #                 size="lg"
        #             )
                
        #         # Right Column
        #         with gr.Column(scale=1):
        #             i2i_output = gr.Image(
        #                 label="Transformed Image",
        #                 type="pil"
        #             )
        #             i2i_status = gr.Textbox(
        #                 label="Status",
        #                 interactive=False
        #             )
        #             i2i_download = gr.File(
        #                 label="Download Image (PNG)",
        #                 file_types=[".png"]
        #             )
            
        #     # Event Handlers
        #     i2i_generate.click(
        #         fn=generate_img2img,
        #         inputs=[i2i_prompt, i2i_input, i2i_strength, i2i_negative, i2i_steps, i2i_cfg, i2i_seed],
        #         outputs=[i2i_output, i2i_status, i2i_download]
        #     )
        
        # ========================================
        # Tab 3: Prompt Helper
        # ========================================
        with gr.Tab("💡 Prompt Helper"):
            gr.Markdown("""
            ## 🎯 Prompt Optimization Guide
            
            **Good prompts are specific and detailed.** Use these templates to improve your results:
            
            ### 📋 Available Styles
            - **Quality**: Best for general high-quality images
            - **Cinematic**: Movie-like scenes with dramatic lighting
            - **Anime**: Japanese animation style
            - **Digital Art**: Concept art and illustrations
            - **Photorealistic**: Ultra-realistic photography
            
            ### 🏷️ Useful Keywords
            
            **Quality Tags**: `masterpiece, best quality, highly detailed, 8k uhd, sharp focus`
            
            **Lighting**: `cinematic lighting, golden hour, volumetric lighting, soft shadows, rim light`
            
            **Camera**: `dslr, 35mm film, wide angle, bokeh, depth of field, macro lens`
            
            **Style**: `oil painting, watercolor, digital art, 3d render, pencil sketch, line art`
            
            **Negative Prompts**: `blur, low quality, distorted, watermark, signature, ugly, duplicate, morbid`
            """)
            
            with gr.Row():
                helper_input = gr.Textbox(
                    label="Your Simple Prompt",
                    placeholder="A cat sitting on a windowsill",
                    lines=2
                )
                helper_style = gr.Dropdown(
                    choices=list(PROMPT_TEMPLATES.keys()),
                    value="quality",
                    label="Style"
                )
            
            helper_output = gr.Textbox(
                label="Optimized Prompt",
                lines=3,
                interactive=False
            )
            
            helper_btn = gr.Button("✨ Generate Optimized Prompt", variant="primary")
            
            helper_btn.click(
                fn=optimize_prompt,
                inputs=[helper_input, helper_style],
                outputs=helper_output
            )
    
    # Footer
    gr.Markdown("""
    ---
    Made with ❤️ using [Gradio](https://gradio.app) & [Hugging Face](https://huggingface.co)
    """)

# ============================================
# Launch Configuration
# ============================================

if __name__ == "__main__":
    demo.queue(max_size=10).launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        favicon_path=None,
        css=CUSTOM_CSS
    )