import gradio as gr
import requests
import os
import time
from PIL import Image
import io
import base64
import json

# Configuration
MIRAGIC_API_URL = os.getenv("MIRAGIC_API_URL")
COMPANY_NAME = "Miragic"
COMPANY_URL = "https://www.miragic.ai"
CONTACT_EMAIL = "info@miragic.ai"

def generate_image_via_api(prompt, model, width, height, seed, image_url, enhance, safe):
    """
    Call the Flask API to generate an image
    """
    try:
        # Prepare the request data
        data = {
            "prompt": prompt,
            "model": model,
            "width": width,
            "height": height,
            "enhance": enhance,
            "safe": safe
        }
        
        if seed and seed != -1:
            data["seed"] = seed
        
        if image_url:
            data["image_url"] = image_url
        
        # Validate kontext model requirements
        if model.lower() == "kontext" and not image_url:
            raise gr.Error("Image URL is required for kontext model (image-to-image generation)")
        
        # Make the API request
        response = requests.post(
            f"{MIRAGIC_API_URL}/generate_image",
            json=data,
            timeout=300  # Increased timeout for image generation
        )
        
        if response.status_code == 200:
            result = response.json()
            if result["status"] == "success":
                # Download the image from Dropbox link
                image_url = result["link"]
                image_response = requests.get(image_url, timeout=30)
                
                if image_response.status_code == 200:
                    # Convert to PIL Image
                    image = Image.open(io.BytesIO(image_response.content))
                    return image, gr.update(visible=True)
                else:
                    return None, gr.update(visible=False)
            else:
                raise gr.Error(f"API Error: {result.get('error', 'Unknown error')}")
        else:
            raise gr.Error(f"HTTP Error: {response.status_code} - {response.text}")
            
    except requests.exceptions.RequestException as e:
        raise gr.Error(f"Request Error: {str(e)}")
    except Exception as e:
        raise gr.Error(f"Unexpected Error: {str(e)}")

def get_available_models():
    """
    Get available models from the Flask API
    """
    try:
        response = requests.get(f"{MIRAGIC_API_URL}/available_models", timeout=10)
        if response.status_code == 200:
            result = response.json()
            if result["status"] == "success":
                return result["models"]
        return ["flux", "kontext", "dreamshaper", "realistic", "anime"]
    except:
        return ["flux", "kontext", "dreamshaper", "realistic", "anime"]

def handle_like_action():
    """Handle like button click"""
    return gr.update(
        value="❤️ Thanks for liking! Please star our repo!",
        interactive=False,
        variant="secondary"
    )

def create_footer_html():
    """Create footer HTML"""
    return f"""
    <div style="text-align: center; margin-top: 30px; padding: 15px; border-radius: 8px;">
        <p style="margin: 5px 0;">© 2025 {COMPANY_NAME}. All rights reserved.</p>
        <p style="margin: 5px 0;">
            <a href="{COMPANY_URL}" target="_blank" style="color: #0984e3; text-decoration: none;">Website</a> | 
            <a href="mailto:{CONTACT_EMAIL}" style="color: #0984e3; text-decoration: none;">Contact Us</a> | 
            <a href="{COMPANY_URL}/privacy-policy" target="_blank" style="color: #0984e3; text-decoration: none;">Privacy Policy</a>
        </p>
    </div>
    """

def create_instructions_html():
    """Create instructions HTML"""
    return """
    <div style="padding: 15px; border-radius: 8px; margin-bottom: 20px;">
        <h2 style="margin-top: 0;">How to use:</h2>
        <ol style="color: #636e72;">
            <li>Enter a detailed description of the image you want to generate</li>
            <li>Select a model (different models have different styles)</li>
            <li>For 'kontext' model, provide an image URL for image-to-image generation</li>
            <li>Adjust the image dimensions and other parameters</li>
            <li>Click "Generate Image" to create your AI artwork</li>
        </ol>
        <p style="margin-bottom: 0;"><strong>Tip:</strong> Be specific in your prompts for better results!</p>
    </div>
    """

def get_css_styles():
    """Get CSS styles for the interface"""
    return """
footer {visibility: hidden}
.banner {
    background-color: #f8f9fa;
    padding: 10px;
    border-radius: 5px;
    margin-bottom: 20px;
    text-align: center;
}
.button-gradient {
  background: linear-gradient(45deg, #ff416c, #ff4b2b, #ff9b00, #ff416c);
  background-size: 400% 400%;
  border: none;
  padding: 14px 28px;
  font-size: 16px;
  font-weight: bold;
  color: white;
  border-radius: 10px;
  cursor: pointer;
  transition: 0.3s ease-in-out;
  animation: gradientAnimation 2s infinite linear;
  box-shadow: 0 4px 10px rgba(255, 65, 108, 0.6);
}
@keyframes gradientAnimation {
  0% { background-position: 0% 50%; }
  100% { background-position: 100% 50%; }
}
.button-gradient:hover {
  transform: scale(1.05);
  box-shadow: 0 6px 15px rgba(255, 75, 43, 0.8);
}
.signup-container {
    text-align: center;
    padding: 20px;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 8px;
    margin-top: 10px;
    color: white;
    box-shadow: 0 4px 15px rgba(0,0,0,0.2);
}
.signup-container h3 {
    margin-bottom: 10px;
    color: white;
}
.signup-container p {
    margin-bottom: 15px;
    color: #f0f0f0;
}
.signup-button {
  background: linear-gradient(45deg, #ff416c, #ff4b2b);
  border: none;
  padding: 12px 25px;
    font-size: 16px;
    font-weight: bold;
    color: white;
    border-radius: 8px;
    text-decoration: none;
    display: inline-block;
    transition: all 0.3s ease;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
}
.signup-button:hover {
    transform: translateY(-2px);
    box-shadow: 0 6px 20px rgba(0, 0, 0, 0.3);
}
#col-container {
    max-width: 1400px;
    margin: 0 auto;
}
.step-column {
    padding: 20px;
    border-radius: 8px;
    box-shadow: var(--card-shadow);
    margin: 10px;
}
.step-title {
    color: var(--primary-color);
    text-align: center;
    margin-bottom: 15px;
}
.image-preview {
    border-radius: 8px;
    overflow: hidden;
    box-shadow: var(--card-shadow);
}
.result-section {
    background-color: white;
    padding: 20px;
    border-radius: 8px;
    box-shadow: var(--card-shadow);
}
.footer {
    margin-top: 30px;
    text-align: center;
}
.like-button {
    background: linear-gradient(45deg, #ff6b6b, #ff8e8e);
    border: none;
    padding: 10px 20px;
    font-size: 14px;
    font-weight: bold;
    color: white;
    border-radius: 8px;
    cursor: pointer;
    transition: all 0.3s ease;
    box-shadow: 0 2px 8px rgba(255, 107, 107, 0.3);
}
.like-button:hover {
    transform: translateY(-2px);
    box-shadow: 0 4px 12px rgba(255, 107, 107, 0.5);
}
.interaction-section {
    text-align: center;
    padding: 20px;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 8px;
    margin-top: 20px;
    color: white;
    box-shadow: 0 4px 15px rgba(0,0,0,0.2);
}
.kontext-notice {
    background: #fff3cd;
    padding: 10px;
    border-radius: 5px;
    margin: 10px 0;
    border-left: 4px solid #ffc107;
    color: #856404;
}
"""

# Create the Gradio interface
with gr.Blocks(css=get_css_styles(), title=f"{COMPANY_NAME} AI Image Generator", theme=gr.themes.Ocean()) as demo:
    with gr.Row():
        with gr.Column(0.7):
            gr.Markdown(f"""
                <div style="display: flex; align-items: center;">
                    <img src="https://avatars.githubusercontent.com/u/211682198?s=200&v=4" style="width: 80px; margin-right: 20px;"/>
                    <div>
                        <h1 style="margin-bottom: 0;">{COMPANY_NAME} AI Image Generator 🎨</h1>
                        <p style="margin-top: 0; color: #636e72;">Create stunning AI-generated images with text prompts</p>
                    </div>
                </div>
                """)

            gr.HTML(create_instructions_html())
        with gr.Column(0.3):
            gr.Markdown(f"""
                <div style="display: flex; align-items: center;">
                    <img src="https://miragic.ai/products/qrcode-vto.png" style="width: 450px; margin-left: 100px;"/>
                </div>
                """)
    
    with gr.Row(elem_id="col-container"):
        # Step 1: Prompt and Settings
        with gr.Column(elem_classes="step-column"):
            gr.HTML("""
            <div class="step-title">
                <span style="font-size: 24px;">1. Enter Prompt & Settings</span><br>
            </div>
            """)
            
            prompt = gr.Textbox(
                label="Prompt", 
                lines=3,
                placeholder="Describe the image you want to generate...",
                value=""
            )
            
            with gr.Row():
                model = gr.Dropdown(
                    choices=get_available_models(),
                    value="flux",
                    label="Model"
                )
                seed = gr.Number(
                    label="Seed", 
                    value=-1,
                    info="-1 for random seed"
                )
            
            with gr.Row():
                width = gr.Slider(
                    minimum=256, 
                    maximum=2048, 
                    value=1024, 
                    step=64, 
                    label="Width"
                )
                height = gr.Slider(
                    minimum=256, 
                    maximum=2048, 
                    value=1024, 
                    step=64, 
                    label="Height"
                )
            
            image_url = gr.Textbox(
                label="Image URL (for kontext model only)",
                placeholder="https://example.com/image.jpg",
                info="Required for image-to-image generation with kontext model"
            )
            
            with gr.Row():
                enhance = gr.Checkbox(label="Enhance Prompt", value=False, info="Use AI to enhance your prompt")
                safe = gr.Checkbox(label="Safe Filter", value=False, info="Enable strict NSFW filtering")
            
            gr.HTML("""
            <div class="kontext-notice">
                <strong>Note:</strong> For the 'kontext' model, you must provide an Image URL for image-to-image generation.
                Example: "bake_a_cake_from_this_logo" with image URL of a logo.
            </div>
            """)
            
            generate_btn = gr.Button(
                "Generate Image 🎨",
                elem_classes="button-gradient"
            )
        
        # Step 2: Results
        with gr.Column(elem_classes="step-column"):
            gr.HTML("""
            <div class="step-title">
                <span style="font-size: 24px;">2. Generated Image</span><br>
            </div>
            """)
            
            result_img = gr.Image(
                label="Generated Image", 
                interactive=False,
                elem_classes="image-preview",
                height=400
            )
            
            with gr.Row():
                like_button = gr.Button(
                    "👍 Like this result!", 
                    elem_classes="like-button",
                    visible=False
                )
            
            gr.HTML("""
            <div class="interaction-section">
                <p style="margin: 5px 0;">If you like our AI Image Generator, please give us a ❤️ into our space!</p>
            </div>
            """)
    
            signup_prompt = gr.HTML(
                visible=True,
                value=f"""<div class="signup-container">
                    <h3>🚀 Want more AI tools?</h3>
                    <p>Visit {COMPANY_NAME}.ai for unlimited access to all our AI tools!</p>
                    <a href='{COMPANY_URL}/' target='_blank' class="signup-button">
                        Explore More Tools 🚀
                    </a>
                </div>"""
            )
    
    # Examples
    examples = [
        ["An epic fantasy landscape with floating islands, cascading waterfalls, ancient ruins, magical aurora borealis in the sky, digital painting, concept art, unreal engine 5, 4K wallpaper", "flux", 2048, 1024, -1, "", False, False],
        ["A serene Japanese garden with cherry blossoms, koi pond, traditional pagoda, morning mist, sunlight filtering through trees, peaceful atmosphere, ultra detailed", "flux", 2048, 1024, -1, "", True, False],
        ["bake_a_cake_from_this_logo", "kontext", 1024, 1024, -1, "https://avatars.githubusercontent.com/u/86964862", False, False],
        ["A minimalist modern living room with large windows, Scandinavian design, plants, cozy atmosphere, architectural rendering", "turbo", 1024, 1024, 42, "", False, True],
        ["A beautiful tropical beach with turquoise water, white sand, palm trees, sunset colors, relaxing vibe, 4K wallpaper", "turbo", 2048, 1024, -1, "", True, False],
    ]
    
    gr.Examples(
        examples=examples,
        inputs=[prompt, model, width, height, seed, image_url, enhance, safe],
        outputs=[result_img, like_button],
        fn=generate_image_via_api,
        cache_examples=False,
        label="Example Prompts"
    )
    
    # Button actions
    generate_btn.click(
        fn=generate_image_via_api,
        inputs=[prompt, model, width, height, seed, image_url, enhance, safe],
        outputs=[result_img, like_button],
        api_name="generate_image"
    )

    like_button.click(
        fn=handle_like_action,
        inputs=[],
        outputs=[like_button]
    )

    # Visitor badge
    gr.HTML(f'<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2F{COMPANY_NAME}-AI%2F{COMPANY_NAME}-AI-Image-Generator"><img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2F{COMPANY_NAME}-AI%2F{COMPANY_NAME}-AI-Image-Generator&label=VISITORS&labelColor=%2337d67a&countColor=%23f47373&style=plastic&labelStyle=upper" /></a>')

    # Footer
    gr.HTML(create_footer_html())
    
if __name__ == "__main__":
    demo.launch(
        show_api=False,
        server_name="0.0.0.0",
        server_port=7860
    )