import spaces
import gradio as gr
import numpy as np
import PIL.Image
from PIL import Image
import random
from diffusers import ControlNetModel, StableDiffusionXLPipeline, AutoencoderKL
from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler
import torch
import os
import time
import glob

# 一時ファイルの保存ディレクトリ
TEMP_DIR = "temp_images"
# 一時ファイルの保持期間（秒）
FILE_RETENTION_PERIOD = 3600  # 1時間

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 一時ディレクトリの作成
os.makedirs(TEMP_DIR, exist_ok=True)

def cleanup_old_files():
    """古い一時ファイルを削除する"""
    current_time = time.time()
    pattern = os.path.join(TEMP_DIR, "output_*.png")
    
    for file_path in glob.glob(pattern):
        try:
            # ファイルの最終更新時刻を取得
            file_modified_time = os.path.getmtime(file_path)
            if current_time - file_modified_time > FILE_RETENTION_PERIOD:
                os.remove(file_path)
        except Exception as e:
            print(f"Error while cleaning up file {file_path}: {e}")

pipe = StableDiffusionXLPipeline.from_single_file(
    "https://huggingface.co/Laxhar/noob_sdxl_beta/noob_hercules3/checkpoint/checkpoint-e2_s109089.safetensors/checkpoint-e2_s109089.safetensors",
    use_safetensors=True,
    torch_dtype=torch.float16,
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1216

@spaces.GPU
def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
    # 古い一時ファイルの削除
    cleanup_old_files()

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)

    # 画像生成
    output_image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator
    ).images[0]

    # RGBモードで保存
    if output_image.mode != 'RGB':
        output_image = output_image.convert('RGB')
    
    # 一時ファイルとして保存
    timestamp = int(time.time())
    temp_filename = os.path.join(TEMP_DIR, f"output_{timestamp}.png")
    output_image.save(temp_filename)
    
    return temp_filename

css = """
#col-container {
    margin: 0 auto;
    width: 100%;
    max-width: 1200px;
    padding: 0 1rem;
}

/* デスクトップレイアウト用のグリッド */
.desktop-layout {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 2rem;
    align-items: start;
}

/* プロンプト入力エリア */
.prompt-container {
    display: flex;
    flex-direction: column;
    gap: 1rem;
}

.prompt-input {
    min-height: 100px !important;
    font-size: 16px !important;
    line-height: 1.5 !important;
    padding: 12px !important;
    border-radius: 8px !important;
    border: 1px solid #e0e0e0 !important;
    background-color: #ffffff !important;
    resize: vertical !important;
}

.prompt-input:focus {
    border-color: #2196f3 !important;
    box-shadow: 0 0 0 2px rgba(33, 150, 243, 0.1) !important;
}

/* 生成ボタン */
.generate-button {
    padding: 12px 24px !important;
    font-size: 16px !important;
    font-weight: 600 !important;
    border-radius: 8px !important;
    background-color: #2196f3 !important;
    color: white !important;
    transition: all 0.3s ease !important;
    margin: 1rem 0 !important;
}

.generate-button:hover {
    background-color: #1976d2 !important;
    transform: translateY(-1px) !important;
}

/* 結果画像 */
#output_image {
    border-radius: 8px;
    overflow: hidden;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
}

/* アコーディオン */
.advanced-settings {
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    overflow: hidden;
    margin-top: 1rem;
}

/* スマートフォン対応 - 768px以下の画面 */
@media (max-width: 768px) {
    .desktop-layout {
        display: block;
    }
    
    #col-container {
        padding: 0 0.5rem;
    }
    
    .prompt-input {
        font-size: 16px !important;
    }
    
    .advanced-settings {
        margin-top: 1rem;
    }
}

/* タブレット対応 - 768px以上1024px以下の画面 */
@media (min-width: 769px) and (max-width: 1024px) {
    .desktop-layout {
        gap: 1rem;
    }
    
    #col-container {
        max-width: 900px;
    }
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
        # Text-to-Image Demo
        Using [Noob SDXL beta model](https://huggingface.co/Laxhar) to generate amazing images!
        """)
        
        with gr.Column(elem_classes="desktop-layout"):
            # 左カラム - 入力コントロール
            with gr.Column(elem_classes="prompt-container"):
                prompt = gr.Textbox(
                    label="What would you like to create?",
                    elem_classes="prompt-input",
                    lines=3,
                    placeholder="Describe the image you want to generate. Be specific about details, style, and atmosphere.\n\nExample: 'A serene mountain landscape at sunset, with snow-capped peaks and a clear lake reflection, painted in watercolor style'",
                    show_label=True,
                )
                run_button = gr.Button(
                    "✨ Generate Image",
                    elem_classes="generate-button",
                    variant="primary",
                )
                
                with gr.Accordion("Advanced Settings", open=False, elem_classes="advanced-settings"):
                    negative_prompt = gr.Textbox(
                        label="Negative Prompt",
                        lines=2,
                        placeholder="Specify what you don't want in the image",
                        value="nsfw, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn"
                    )

                    with gr.Row():
                        with gr.Column(scale=3):
                            seed = gr.Slider(
                                label="Seed",
                                minimum=0,
                                maximum=MAX_SEED,
                                step=1,
                                value=0,
                            )
                        with gr.Column(scale=1):
                            randomize_seed = gr.Checkbox(
                                label="Randomize",
                                value=True,
                            )

                    with gr.Row():
                        width = gr.Slider(
                            label="Width",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=32,
                            value=1024,
                        )
                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=32,
                            value=1024,
                        )

                    with gr.Row():
                        guidance_scale = gr.Slider(
                            label="Guidance Scale",
                            minimum=0.0,
                            maximum=20.0,
                            step=0.1,
                            value=7,
                            info="Controls how closely the image follows the prompt"
                        )
                        num_inference_steps = gr.Slider(
                            label="Steps",
                            minimum=1,
                            maximum=28,
                            step=1,
                            value=28,
                            info="More steps = higher quality"
                        )



            # 右カラム - 生成結果
            with gr.Column():
                result = gr.Image(
                    label="Generated Image",
                    show_label=True,
                    type="filepath",
                    elem_id="output_image"
                )

    run_button.click(
        fn=infer,
        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[result]
    )

# 起動時に古いファイルを削除
cleanup_old_files()

demo.queue().launch()