import os
import random
import sys
from typing import Sequence, Mapping, Any, Union
import torch
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download
import base64
from io import BytesIO
from openai import OpenAI

def init_gpt_api():
    return OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def base64_to_pil(base64_str):
    image_data = base64.b64decode(base64_str)
    return Image.open(BytesIO(image_data))
    
def customize_image(prompt, image):
    if image is None or prompt.strip() == "":
        return None

    prompt = f"Convert this image into a {prompt} drawing"
    client = init_gpt_api()
    # image_base64 = encode_image_from_pil(image)

    result = client.images.edit(
        model="gpt-image-1",
        image=[open(image, 'rb')],  
        prompt=prompt,
    )
    if result.data:
        image_base64 = result.data[0].b64_json
        return base64_to_pil(image_base64)
    else:
        return None
        

# Create Gradio interface

custom_css = """
/* Global styling */
.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    min-height: 100vh;
    padding: 20px;
}

/* Header styling */
.main-header {
    background: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(10px);
    border-radius: 20px;
    padding: 40px;
    margin-bottom: 30px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
    border: 1px solid rgba(255, 255, 255, 0.2);
    text-align: center;
}

.main-header h1 {
    background: linear-gradient(45deg, #667eea, #764ba2);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-size: 3rem !important;
    font-weight: 700 !important;
    margin-bottom: 15px !important;
    text-align: center !important;
}

.main-header h3 {
    color: #6b7280 !important;
    text-align: center !important;
    font-weight: 400 !important;
    font-size: 1.3rem !important;
    line-height: 1.6 !important;
}

/* Step containers */
.step-container {
    background: rgba(255, 255, 255, 0.9);
    backdrop-filter: blur(10px);
    border-radius: 16px;
    padding: 30px;
    margin-bottom: 25px;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.08);
    border: 1px solid rgba(255, 255, 255, 0.3);
    transition: transform 0.2s ease, box-shadow 0.2s ease;
}

.step-container:hover {
    transform: translateY(-2px);
    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.12);
}

/* Step headers */
.step-header {
    display: flex;
    align-items: center;
    margin-bottom: 25px;
    padding-bottom: 20px;
    border-bottom: 3px solid #f3f4f6;
}

.step-number {
    background: linear-gradient(45deg, #667eea, #764ba2);
    color: white;
    width: 45px;
    height: 45px;
    border-radius: 50%;
    display: flex;
    align-items: center;
    justify-content: center;
    font-weight: bold;
    font-size: 1.3rem;
    margin-right: 20px;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3);
}

.step-title {
    font-size: 1.5rem;
    font-weight: 600;
    color: #374151;
    margin: 0;
}

/* Enhanced button styling */
.primary-button {
    background: linear-gradient(45deg, #667eea, #764ba2) !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 15px 35px !important;
    font-weight: 600 !important;
    font-size: 1.1rem !important;
    color: white !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3) !important;
    transition: all 0.3s ease !important;
    margin: 15px 5px !important;
    min-height: 50px !important;
    width: 100% !important;
}

.primary-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 25px rgba(102, 126, 234, 0.4) !important;
}

.secondary-button {
    background: linear-gradient(45deg, #10b981, #059669) !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 15px 35px !important;
    font-weight: 600 !important;
    font-size: 1.1rem !important;
    color: white !important;
    box-shadow: 0 4px 15px rgba(16, 185, 129, 0.3) !important;
    transition: all 0.3s ease !important;
    margin: 15px 5px !important;
    min-height: 50px !important;
    width: 100% !important;
}

.secondary-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 25px rgba(16, 185, 129, 0.4) !important;
}

/* Enhanced input styling */
.gr-textbox input, .gr-dropdown, .gr-slider {
    border: 2px solid #e5e7eb !important;
    border-radius: 10px !important;
    padding: 12px 16px !important;
    font-size: 1rem !important;
    transition: all 0.3s ease !important;
}

.gr-textbox input:focus, .gr-dropdown:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
    outline: none !important;
}

/* Upload area enhancement */
.upload-area {
    border: 3px dashed #667eea;
    border-radius: 15px;
    padding: 10px;
    text-align: center;
    background: linear-gradient(135deg, rgba(102, 126, 234, 0.05), rgba(118, 75, 162, 0.05));
    transition: all 0.3s ease;
    min-height: 150px;
    display: flex;
    align-items: center;
    justify-content: center;
}

.upload-area:hover {
    border-color: #764ba2;
    background: linear-gradient(135deg, rgba(102, 126, 234, 0.1), rgba(118, 75, 162, 0.1));
    transform: translateY(-2px);
}

/* Image preview styling */
.image-preview {
    border: 2px solid #e2e8f0;
    border-radius: 12px;
    overflow: hidden;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
    background: #f8fafc;
}

/* Gallery styling */
.enhanced-gallery {
    border: 2px solid #e2e8f0;
    border-radius: 12px;
    padding: 15px;
    background: #f8fafc;
    min-height: 300px;
}

/* Control groups */
.control-group {
    background: #f8fafc;
    border: 2px solid #e2e8f0;
    border-radius: 12px;
    padding: 20px;
    margin: 15px 0;
}

.control-group h4 {
    color: #374151;
    font-weight: 600;
    margin-bottom: 15px;
    text-align: center;
}

/* Info boxes */
.info-box {
    background: linear-gradient(45deg, #dbeafe, #bfdbfe);
    border: 2px solid #60a5fa;
    border-radius: 10px;
    padding: 15px 20px;
    margin: 15px 0;
    color: #1e40af;
    font-weight: 500;
}

.warning-box {
    background: linear-gradient(45deg, #fef3c7, #fde68a);
    border: 2px solid #f59e0b;
    border-radius: 10px;
    padding: 15px 20px;
    margin: 15px 0;
    color: #92400e;
    font-weight: 500;
}

.success-box {
    background: linear-gradient(45deg, #d1fae5, #a7f3d0);
    border: 2px solid #10b981;
    border-radius: 10px;
    padding: 15px 20px;
    margin: 15px 0;
    color: #065f46;
    font-weight: 500;
}

/* Section dividers */
.section-divider {
    height: 3px;
    background: linear-gradient(45deg, #667eea, #764ba2);
    border-radius: 2px;
    margin: 30px 0;
    opacity: 0.7;
}

/* Examples section */
.examples-section {
    background: rgba(255, 255, 255, 0.9);
    backdrop-filter: blur(10px);
    border-radius: 16px;
    padding: 30px;
    margin-top: 30px;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.08);
}
/* Responsive design */
@media (max-width: 768px) {
    .gradio-container {
        padding: 15px;
    }
    
    .main-header {
        padding: 25px;
    }
    
    .main-header h1 {
        font-size: 2.5rem !important;
    }
    
    .step-container {
        padding: 20px;
    }
    
    .step-number {
        width: 35px;
        height: 35px;
        font-size: 1.1rem;
    }
    
    .step-title {
        font-size: 1.3rem;
    }
}

/* Hide default gradio styling conflicts */
.gr-box {
    margin-bottom: 0px !important;
    border: none !important;
}

.gr-panel {
    border: none !important;
    background: transparent !important;
}

.gr-form {
    background: transparent !important;
}
"""
examples = [
    ["", "mona.png", "receita-tacos.webp", 15, 0.6],
    ["a woman looking at a house catching fire on the background", "disaster_girl.png", "abaporu.jpg", 15, 0.15],
    ["istanbul aerial, dramatic photography", "natasha.png", "istambul.jpg", 15, 0.5],
]
output_image = gr.Image(label="Generated Image")
multi_image = gr.Image(label="Generated Image")
# customized_image = gr.Image(label="Customized Image")
medium_images = gr.Gallery(label="Images", columns=[2], rows=[2], elem_id="gallery", format="png")


with gr.Blocks(css=custom_css, title="🎨 Artist Ideation - Advanced Style Transfer") as app:
    # Enhanced Header
    gr.HTML("""
    <div class="main-header">
        <h1>🎨 Artist Ideation</h1>
        <h3>Advanced Image Medium Transfer </h3>
    </div>
    """)

    with gr.Row(elem_classes="step-container"):
        with gr.Column():
            gr.HTML("""
            <div class="step-header">
                <h2 class="step-title">Artistic Medium Transformation</h2>
            </div>
            """)
            with gr.Row():
                with gr.Column(scale=1):
                    structure_image = gr.Image(label="Structure Image", type="filepath")
                    custom_medium = gr.Textbox(
                        label="🎨 Artistic Medium", 
                        placeholder="Enter your desired medium (e.g., watercolor, charcoal sketch, digital art)...",
                        lines=2,
                    )
                    gr.HTML("""
                    <div class="success-box">
                        💡 <strong>Pro Tips:</strong> Try "watercolor portrait", "pencil sketch", "oil painting masterpiece", 
                        or "digital art concept" for best results!
                    </div>
                    """)
                    customize_btn = gr.Button(
                        "🖌️ Transform Medium", 
                        variant="primary",
                        elem_classes="secondary-button",
                        size="lg"
                    )
                    
                with gr.Column(scale=1):
                    customized_image = gr.Image(
                        label="🎭 Artistic Transformation", 
                        elem_classes="image-preview",
                        height=400
                    )
                    # customized_image.render()
                    
    customize_btn.click(
        fn=customize_image,
        inputs=[custom_medium, structure_image],
        outputs=[customized_image]
    )

if __name__ == "__main__":
    app.launch(share=True)