import os
import gc
import gradio as gr
import numpy as np
import torch
import json
import spaces
import random
import config
import utils
import logging
import prompt_generator
from PIL import Image, PngImagePlugin
from datetime import datetime
from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
from config import (
    MODEL,
    MIN_IMAGE_SIZE,
    MAX_IMAGE_SIZE,
    USE_TORCH_COMPILE,
    ENABLE_CPU_OFFLOAD,
    OUTPUT_DIR,
    DEFAULT_NEGATIVE_PROMPT,
    DEFAULT_ASPECT_RATIO,
    sampler_list,
    aspect_ratios,
    style_list,
    # 設定
    TEXT_TO_PROMPT_ENABLED,
    DEFAULT_CATEGORY,
    DEFAULT_SERIES,
    DEFAULT_CHARACTER,
    series_list,
    character_list,
    category_list,
)
import time
from typing import List, Dict, Tuple, Optional

# Enhanced logging configuration
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# Constants
IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
HF_TOKEN = os.getenv("HF_TOKEN")
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"

# Create CSS with improved buttons and styling
custom_css = """
.header {
    text-align: center;
    margin-bottom: 2rem;
    background: linear-gradient(to right, #4a69bd, #6a89cc);
    padding: 1.5rem;
    border-radius: 10px;
    color: white;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}

.title {
    margin: 0;
    font-size: 2.5rem;
    font-weight: 700;
}

.subtitle {
    font-size: 1.1rem;
    margin-top: 0.5rem;
    opacity: 0.9;
}

.subtitle-inline {
    font-size: 1.3rem;
    font-weight: 400;
    opacity: 0.9;
}

.notification {
    background-color: #fff8e1;
    border-left: 5px solid #ffc107;
    padding: 20px;
    margin: 20px 0;
    font-size: 1.2rem;
    border-radius: 10px;
    box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);
}

.notification-title {
    color: #e65100;
    font-size: 1.5rem;
    margin-bottom: 10px;
    font-weight: 600;
}

.en-message {
    margin-bottom: 15px;
}

.jp-message {
    font-weight: 500;
}
"""

# Create the Gradio interface
with gr.Blocks(css=custom_css) as demo:
    gr.HTML("<div class='header'><h1 class='title'>FanFic Illustrator <span class='subtitle-inline'>with Animagine XL 4.0 Opt</span></h1><p class='subtitle'>Illustrate your fan stories with beautiful AI-generated art<br>二次創作ファン小説にAIで魅力的な挿絵を</p></div>")
    
    with gr.Column():
        # Service temporarily unavailable notification
        gr.HTML("""
        <div class="notification">
            <div class="notification-title">Service Temporarily Suspended</div>
            <div class="en-message">
                This service has been temporarily suspended due to frequent ZERO GPU allocation failures.<br>
                You can try it for free using <a href="https://github.com/webbigdata-jp/python_sample/blob/main/FanFic_Illustrator_demo.ipynb">Google Colab FanFic Illustrator demo</a>
            </div>
            <div class="jp-message">
                ZERO GPUの割当に失敗する事が多すぎるので一時停止しました。<br>
                <a href="https://github.com/webbigdata-jp/python_sample/blob/main/FanFic_Illustrator_demo.ipynb">Google Colab FanFic Illustrator demo</a>で無料で試す事ができます
            </div>
        </div>
        """)

# Launch the app
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", share=IS_COLAB)