import os
import gradio as gr
import base64
from random import randint
from all_models import models
from io import BytesIO
from PIL import Image
from fastapi import FastAPI, Request
from deep_translator import GoogleTranslator

# CSS yang lebih lengkap dengan berbagai elemen styling
css_code = """
/* General Styling */
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    min-height: 100vh;
    padding: 20px;
}

/* Textbox Styling */
#custom_textbox {
    min-height: 150px;
    padding: 15px;
    border-radius: 12px;
    border: 2px solid #e0e0e0;
    font-size: 16px;
    background: rgba(255, 255, 255, 0.9);
    box-shadow: 0 2px 10px rgba(0, 0, 0, 0.05);
    transition: all 0.3s ease;
}

#custom_textbox:focus {
    border-color: #4facfe;
    box-shadow: 0 0 0 3px rgba(79, 172, 254, 0.2);
    outline: none;
}

/* Button Styling */
#custom_gen_button {
    background: linear-gradient(to right, #4facfe 0%, #00f2fe 100%);
    color: white;
    border: none;
    border-radius: 12px;
    padding: 12px 24px;
    font-weight: 600;
    font-size: 16px;
    cursor: pointer;
    transition: all 0.3s ease;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    width: 100%;
    margin-bottom: 10px;
}

#custom_gen_button:hover {
    transform: translateY(-2px);
    box-shadow: 0 6px 12px rgba(0, 0, 0, 0.15);
}

#custom_stop_button {
    background: linear-gradient(to right, #ff4d4d 0%, #f97878 100%);
    color: white;
    border: none;
    border-radius: 12px;
    padding: 12px 24px;
    font-weight: 600;
    font-size: 16px;
    cursor: pointer;
    transition: all 0.3s ease;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    width: 100%;
}

#custom_stop_button:hover {
    transform: translateY(-2px);
    box-shadow: 0 6px 12px rgba(0, 0, 0, 0.15);
}

/* Image Styling */
#custom_image {
    border-radius: 16px;
    box-shadow: 0 8px 24px rgba(0, 0, 0, 0.1);
    transition: all 0.3s ease;
    background: white;
    padding: 10px;
    max-width: 100%;
    height: auto;
}

#custom_image:hover {
    transform: scale(1.01);
    box-shadow: 0 12px 28px rgba(0, 0, 0, 0.15);
}

/* Dropdown Styling */
.dark .dropdown-option {
    background: #2d3748 !important;
}

.dark .dropdown-option:hover {
    background: #4a5568 !important;
}

/* Panel Styling */
.gr-box {
    border-radius: 16px !important;
    background: rgba(255, 255, 255, 0.8) !important;
    backdrop-filter: blur(10px);
    border: 1px solid rgba(255, 255, 255, 0.2) !important;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1) !important;
    padding: 20px !important;
}

/* Label Styling */
.gr-label {
    font-weight: 600 !important;
    color: #2d3748 !important;
    margin-bottom: 8px !important;
    font-size: 16px !important;
}

/* Responsive Adjustments */
@media (max-width: 768px) {
    .gradio-container {
        padding: 10px;
    }
    
    #custom_textbox {
        min-height: 120px;
        font-size: 14px;
    }
    
    #custom_gen_button, #custom_stop_button {
        padding: 10px 15px;
        font-size: 14px;
    }
}

/* Animation for loading state */
@keyframes pulse {
    0% { opacity: 0.6; }
    50% { opacity: 1; }
    100% { opacity: 0.6; }
}

.loading {
    animation: pulse 1.5s infinite;
}
"""

# Initialize translator
translator = GoogleTranslator(source='auto', target='en')

# Load models
models_load = {}
for model in models:
    try:
        models_load[model] = gr.load(f'models/{model}')
    except Exception as error:
        models_load[model] = gr.Interface(lambda txt: None, ['text'], ['image'])

app = FastAPI()

def gen_image(model_str, prompt):
    if model_str == 'NA':
        return None
    # Translate prompt to English {noise} {klir}
    translated_prompt = translator.translate(prompt)
    noise = str(randint(0, 4294967296))
    klir = '| ultra detail, ultra elaboration, ultra quality, perfect'
    return models_load[model_str](f'{translated_prompt}')

def image_to_base64(image):
    buffered = BytesIO()
    if isinstance(image, str):  # if it's a file path
        img = Image.open(image)
        img.save(buffered, format="JPEG")
    else:  # if it's a PIL Image
        image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode()

# API endpoint
@app.post("/generate")
async def api_generate(request: Request):
    data = await request.json()
    model = data.get('model', models[0])
    prompt = data.get('prompt', '')
    
    if model not in models:
        return {"error": "Model not found"}
    
    # Translate prompt to English for API endpoint too
    translated_prompt = translator.translate(prompt)
    image = gen_image(model, translated_prompt)
    if image is None:
        return {"error": "Image generation failed"}
    
    base64_str = image_to_base64(image)
    
    return {
        "status": "success",
        "model": model,
        "original_prompt": prompt,
        "translated_prompt": translated_prompt,
        "image_base64": base64_str,
        "image_format": "jpeg"
    }

# Gradio Interface
def make_me():
    with gr.Row():
        with gr.Column(scale=4):
            txt_input = gr.Textbox(
                label='Your prompt:', 
                lines=4, 
                container=False, 
                elem_id="custom_textbox", 
                placeholder="Describe the image you want to generate...",
                interactive=True
            )
        
        with gr.Column(scale=1):
            gen_button = gr.Button('Generate image', elem_id="custom_gen_button", variant="primary")
            stop_button = gr.Button('Stop', variant='secondary', interactive=False, 
                                elem_id="custom_stop_button")
            
            def on_generate_click():
                return gr.Button('Generating...', interactive=False, elem_id="custom_gen_button", variant="stop"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button")
            
            def on_stop_click():
                return gr.Button('Generate image', elem_id="custom_gen_button", variant="primary"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
            
            gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
            stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
    
    with gr.Row():
        with gr.Column():
            model_dropdown = gr.Dropdown(models, label="Select Model", 
                                      value=models[0] if models else None,
                                      interactive=True,
                                      elem_classes=["model-selector"])
            output_image = gr.Image(
                label="Generated Image", 
                width=512,
                height=768,
                elem_id="custom_image", 
                show_label=True, 
                interactive=False
            )
            
            # JSON output
            json_output = gr.JSON(label="API Response", elem_id="api-response")
            
            def generate_wrapper(model_str, prompt):
                # Translate prompt to English
                translated_prompt = translator.translate(prompt)
                image = gen_image(model_str, translated_prompt)
                if image is None:
                    return None, {"error": "Generation failed"}
                
                base64_str = image_to_base64(image)
                response = {
                    "status": "success",
                    "model": model_str,
                    "original_prompt": prompt,
                    "translated_prompt": translated_prompt,
                    "image_base64": base64_str,
                    "image_format": "jpeg"
                }
                return image, response
            
            gen_event = gen_button.click(generate_wrapper, [model_dropdown, txt_input], 
                                     [output_image, json_output])
            stop_button.click(on_stop_click, inputs=None, 
                            outputs=[gen_button, stop_button], cancels=[gen_event])

# Create Gradio app
with gr.Blocks(css=css_code, theme=gr.themes.Soft()) as demo: 
    gr.Markdown("# AI Image Generator", elem_id="title")
    gr.Markdown("Generate stunning images from text prompts using advanced AI models", elem_id="subtitle")
    make_me()

# Enable queue before mounting
demo.queue()

# Mount Gradio app to FastAPI
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)