import os
import tempfile
from PIL import Image, ImageEnhance, ImageFilter
import gradio as gr
import logging
import re
import time
import cv2
import numpy as np
from io import BytesIO
from datetime import datetime, timedelta
import random

from google import genai
from google.genai import types
from dotenv import load_dotenv

load_dotenv()

# 로깅 설정
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# ------------------- API 키 순환 시스템 -------------------
# API 키 설정 - 환경변수에서 API_KEYS = [...] 구문 그대로 실행
def load_api_configs_from_env():
    env_configs = os.getenv('API_KEYS')
    if env_configs:
        try:
            # API_KEYS = [...] 구문을 그대로 실행
            exec(env_configs, globals())
            return globals().get('API_KEYS', [])
        except Exception as e:
            logger.error(f"⚠️ API_KEYS 환경변수 실행 오류: {e}")
            return []
    return []

# 환경변수에서 로드하거나 기본값 사용
loaded_configs = load_api_configs_from_env()
API_KEYS = loaded_configs if loaded_configs else [
    "",
    "",
    "",
    "",
    "",
    "",
    "",
    "",
    "",
    "",
]

current_key_index = -1  # -1은 아직 사용하지 않았음을 의미

def initialize_api_keys():
    """API 키 목록을 초기화하는 함수"""
    global API_KEYS
    
    # 빈 문자열이나 None인 키들 제거
    API_KEYS = [key for key in API_KEYS if key and key.strip() and not key.startswith("your_")]
    
    logger.info(f"API 키 {len(API_KEYS)}개가 로드되었습니다.")
    
    if len(API_KEYS) == 0:
        logger.error("사용 가능한 API 키가 없습니다. API_KEYS 리스트에 실제 키를 입력해주세요.")

def get_next_api_key():
    """다음 API 키를 가져오는 함수 (처음에는 랜덤, 이후에는 순차적)"""
    global current_key_index, API_KEYS
    
    if not API_KEYS:
        logger.error("API 키가 설정되지 않았습니다.")
        return None
    
    if current_key_index == -1:
        # 첫 번째 사용은 랜덤으로 시작점 선택
        current_key_index = random.randint(0, len(API_KEYS) - 1)
        logger.info(f"첫 번째 API 키 랜덤 선택: 인덱스 {current_key_index}")
    else:
        # 이후에는 순차적으로 다음 키 사용
        current_key_index = (current_key_index + 1) % len(API_KEYS)
        logger.info(f"다음 API 키 순차 선택: 인덱스 {current_key_index}")
    
    return API_KEYS[current_key_index]

# ========== 이미지 생성기 관련 함수 ==========
def save_binary_file(file_name, data):
    with open(file_name, "wb") as f:
        f.write(data)

def translate_prompt_to_english(prompt):
    if not re.search("[가-힣]", prompt):
        return prompt
    
    prompt = prompt.replace("#1", "IMAGE_TAG_ONE")
    prompt = prompt.replace("#2", "IMAGE_TAG_TWO")
    prompt = prompt.replace("#3", "IMAGE_TAG_THREE")
    
    try:
        api_key = get_next_api_key()
        if not api_key:
            logger.error("Gemini API 키가 설정되지 않았습니다.")
            prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
            prompt = prompt.replace("IMAGE_TAG_TWO", "#2")
            prompt = prompt.replace("IMAGE_TAG_THREE", "#3")
            return prompt
            
        client = genai.Client(api_key=api_key)
        translation_prompt = f"""
        Translate the following Korean text to English:
        
        {prompt}
        
        IMPORTANT: The tokens IMAGE_TAG_ONE, IMAGE_TAG_TWO, and IMAGE_TAG_THREE are special tags 
        and must be preserved exactly as is in your translation. Do not translate these tokens.
        """
        
        logger.info(f"Translation prompt: {translation_prompt}")
        response = client.models.generate_content(
            model="gemini-2.0-flash",
            contents=[translation_prompt],
            config=types.GenerateContentConfig(
                response_modalities=['Text'],
                temperature=0.2,
                top_p=0.95,
                top_k=40,
                max_output_tokens=512
            )
        )
        
        translated_text = ""
        for part in response.candidates[0].content.parts:
            if hasattr(part, 'text') and part.text:
                translated_text += part.text
        
        if translated_text.strip():
            translated_text = translated_text.replace("IMAGE_TAG_ONE", "#1")
            translated_text = translated_text.replace("IMAGE_TAG_TWO", "#2")
            translated_text = translated_text.replace("IMAGE_TAG_THREE", "#3")
            logger.info(f"Translated text: {translated_text.strip()}")
            return translated_text.strip()
        else:
            logger.warning("번역 결과가 없습니다. 원본 프롬프트 사용")
            prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
            prompt = prompt.replace("IMAGE_TAG_TWO", "#2")
            prompt = prompt.replace("IMAGE_TAG_THREE", "#3")
            return prompt
    except Exception as e:
        logger.exception("번역 중 오류 발생:")
        prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
        prompt = prompt.replace("IMAGE_TAG_TWO", "#2")
        prompt = prompt.replace("IMAGE_TAG_THREE", "#3")
        return prompt

def preprocess_prompt(prompt, image1, image2, image3):
    # 기존 함수 유지
    has_img1 = image1 is not None
    has_img2 = image2 is not None
    has_img3 = image3 is not None

    if "#1" in prompt and not has_img1:
        prompt = prompt.replace("#1", "첫 번째 이미지(없음)")
    else:
        prompt = prompt.replace("#1", "첫 번째 이미지")

    if "#2" in prompt and not has_img2:
        prompt = prompt.replace("#2", "두 번째 이미지(없음)")
    else:
        prompt = prompt.replace("#2", "두 번째 이미지")

    if "#3" in prompt and not has_img3:
        prompt = prompt.replace("#3", "세 번째 이미지(없음)")
    else:
        prompt = prompt.replace("#3", "세 번째 이미지")

    if "1. 이미지 변경" in prompt:
        desc_match = re.search(r'#1을 "(.*?)"으로 바꿔라', prompt)
        if desc_match:
            description = desc_match.group(1)
            prompt = f"첫 번째 이미지를 {description}으로 변경해주세요. 원본 이미지의 주요 내용은 유지하되 새로운 스타일과 분위기로 재해석해주세요."
        else:
            prompt = "첫 번째 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."

    elif "2. 글자지우기" in prompt:
        text_match = re.search(r'#1에서 "(.*?)"를 지워라', prompt)
        if text_match:
            text_to_remove = text_match.group(1)
            prompt = f"첫 번째 이미지에서 '{text_to_remove}' 텍스트를 찾아 자연스럽게 제거해주세요. 텍스트가 있던 부분을 배경과 조화롭게 채워주세요."
        else:
            prompt = "첫 번째 이미지에서 모든 텍스트를 찾아 자연스럽게 제거해주세요. 깔끔한 이미지로 만들어주세요."

    elif "4. 옷바꾸기" in prompt:
        prompt = "첫 번째 이미지의 인물 의상을 두 번째 이미지의 의상으로 변경해주세요. 의상의 스타일과 색상은 두 번째 이미지를 따르되, 신체 비율과 포즈는 첫 번째 이미지를 유지해주세요."

    elif "5. 배경바꾸기" in prompt:
        prompt = "첫 번째 이미지의 배경을 두 번째 이미지의 배경으로 변경해주세요. 첫 번째 이미지의 주요 피사체는 유지하고, 두 번째 이미지의 배경과 조화롭게 합성해주세요."

    elif "6. 이미지 합성(상품포함)" in prompt:
        prompt = "첫 번째 이미지와 두 번째 이미지(또는 세 번째 이미지)를 자연스럽게 합성해주세요. 모든 이미지의 주요 요소를 포함하고, 특히 상품이 돋보이도록 조화롭게 통합해주세요."

    prompt += " 이미지를 생성해주세요. 이미지에 텍스트나 글자를 포함하지 마세요."
    return prompt

def generate_with_images(prompt, images, variation_index=0):
    try:
        api_key = get_next_api_key()
        if not api_key:
            return None, "API 키가 설정되지 않았습니다. API_KEYS 환경변수에 API 키를 설정해주세요."

        client = genai.Client(api_key=api_key)
        logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}, 변형 인덱스: {variation_index}")

        variation_suffixes = [
            " Create this as the first variation. Do not add any text, watermarks, or labels to the image.",
            " Create this as the second variation with more vivid colors. Do not add any text, watermarks, or labels to the image."
        ]
        
        if variation_index < len(variation_suffixes):
            prompt = prompt + variation_suffixes[variation_index]
        else:
            prompt = prompt + " Do not add any text, watermarks, or labels to the image."

        contents = [prompt]
        for idx, img in enumerate(images, 1):
            if img is not None:
                contents.append(img)
                logger.info(f"이미지 #{idx} 추가됨")

        response = client.models.generate_content(
            model="gemini-2.0-flash-exp-image-generation",
            contents=contents,
            config=types.GenerateContentConfig(
                response_modalities=['Text', 'Image'],
                temperature=1,
                top_p=0.95,
                top_k=40,
                max_output_tokens=8192
            )
        )

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            temp_path = tmp.name
            result_text = ""
            image_found = False
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'text') and part.text:
                    result_text += part.text
                    logger.info(f"응답 텍스트: {part.text}")
                elif hasattr(part, 'inline_data') and part.inline_data:
                    save_binary_file(temp_path, part.inline_data.data)
                    image_found = True
                    logger.info("응답에서 이미지 추출 성공")
            if not image_found:
                return None, f"API에서 이미지를 생성하지 못했습니다. 응답 텍스트: {result_text}"
            result_img = Image.open(temp_path)
            if result_img.mode == "RGBA":
                result_img = result_img.convert("RGB")
            return result_img, f"이미지가 성공적으로 생성되었습니다. {result_text}"
    except Exception as e:
        logger.exception("이미지 생성 중 오류 발생:")
        return None, f"오류 발생: {str(e)}"

def process_images_with_prompt(image1, image2, image3, prompt, variation_index=0, max_retries=3):
    retry_count = 0
    last_error = None
    
    while retry_count < max_retries:
        try:
            images = [image1, image2, image3]
            valid_images = [img for img in images if img is not None]
            
            if not valid_images:  # 여기가 254번째 줄 근처일 것
                if prompt and prompt.strip():
                    # 프롬프트만 있는 경우 텍스트 기반 생성 시도
                    logger.info("텍스트 프롬프트만으로 이미지 생성 시도")
                    try:
                        # 빈 이미지 리스트로 생성 시도
                        result_img, status = generate_with_images(prompt, [], variation_index)
                        if result_img is not None:
                            return result_img, status, prompt
                    except:
                        pass
                return None, "이미지를 업로드하거나 텍스트 프롬프트를 입력해주세요.", ""           

            if prompt and prompt.strip():
                processed_prompt = preprocess_prompt(prompt, image1, image2, image3)
                if re.search("[가-힣]", processed_prompt):
                    final_prompt = translate_prompt_to_english(processed_prompt)
                else:
                    final_prompt = processed_prompt
            else:
                if len(valid_images) == 1:
                    final_prompt = "Please creatively transform this image into a more vivid and artistic version. Do not include any text or watermarks in the generated image."
                    logger.info("Default prompt generated for single image")
                elif len(valid_images) == 2:
                    final_prompt = "Please seamlessly composite these two images, integrating their key elements harmoniously into a single image. Do not include any text or watermarks in the generated image."
                    logger.info("Default prompt generated for two images")
                else:
                    final_prompt = "Please creatively composite these three images, combining their main elements into a cohesive and natural scene. Do not include any text or watermarks in the generated image."
                    logger.info("Default prompt generated for three images")

            result_img, status = generate_with_images(final_prompt, valid_images, variation_index)
            if result_img is not None:
                return result_img, status, final_prompt
            else:
                last_error = status
                retry_count += 1
                logger.warning(f"이미지 생성 실패, 재시도 {retry_count}/{max_retries}: {status}")
                time.sleep(1)
        except Exception as e:
            last_error = str(e)
            retry_count += 1
            logger.exception(f"이미지 처리 중 오류 발생, 재시도 {retry_count}/{max_retries}:")
            time.sleep(1)
    
    return None, f"최대 재시도 횟수({max_retries}회) 초과 후 실패: {last_error}", prompt

def generate_multiple_images(image1, image2, image3, prompt, progress=gr.Progress()):
    results = []
    statuses = []
    prompts = []
    
    num_images = 2
    max_retries = 3
    
    progress(0, desc="이미지 생성 준비 중...")
    
    for i in range(num_images):
        progress((i / num_images), desc=f"{i+1}/{num_images} 이미지 생성 중...")
        result_img, status, final_prompt = process_images_with_prompt(image1, image2, image3, prompt, i, max_retries)
        
        if result_img is not None:
            results.append(result_img)
            statuses.append(f"이미지 #{i+1}: {status}")
            prompts.append(f"이미지 #{i+1}: {final_prompt}")
        else:
            results.append(None)
            statuses.append(f"이미지 #{i+1} 생성 실패: {status}")
            prompts.append(f"이미지 #{i+1}: {final_prompt}")
        
        time.sleep(1)
    
    progress(1.0, desc="이미지 생성 완료!")
    
    while len(results) < 2:
        results.append(None)
    
    combined_status = "\n".join(statuses)
    combined_prompts = "\n".join(prompts)
    
    return results[0], results[1], combined_status, combined_prompts

# 다크모드 자동 적용 커스텀 CSS 스타일
custom_css = """
/* ============================================
   다크모드 자동 변경 템플릿 CSS
   ============================================ */
/* 1. CSS 변수 정의 (라이트모드 - 기본값) */
:root {
    /* 메인 컬러 */
    --primary-color: #FB7F0D;
    --secondary-color: #ff9a8b;
    --accent-color: #FF6B6B;
    
    /* 배경 컬러 */
    --background-color: #FFFFFF;
    --card-bg: #ffffff;
    --input-bg: #ffffff;
    
    /* 텍스트 컬러 */
    --text-color: #334155;
    --text-secondary: #64748b;
    
    /* 보더 및 구분선 */
    --border-color: #dddddd;
    --border-light: #e5e5e5;
    
    /* 테이블 컬러 */
    --table-even-bg: #f3f3f3;
    --table-hover-bg: #f0f0f0;
    
    /* 그림자 */
    --shadow: 0 8px 30px rgba(251, 127, 13, 0.08);
    --shadow-light: 0 2px 4px rgba(0, 0, 0, 0.1);
    
    /* 기타 */
    --border-radius: 18px;
}
/* 2. 다크모드 색상 변수 (자동 감지) */
@media (prefers-color-scheme: dark) {
    :root {
        /* 배경 컬러 */
        --background-color: #1a1a1a;
        --card-bg: #2d2d2d;
        --input-bg: #2d2d2d;
        
        /* 텍스트 컬러 */
        --text-color: #e5e5e5;
        --text-secondary: #a1a1aa;
        
        /* 보더 및 구분선 */
        --border-color: #404040;
        --border-light: #525252;
        
        /* 테이블 컬러 */
        --table-even-bg: #333333;
        --table-hover-bg: #404040;
        
        /* 그림자 */
        --shadow: 0 8px 30px rgba(0, 0, 0, 0.3);
        --shadow-light: 0 2px 4px rgba(0, 0, 0, 0.2);
    }
}
/* 3. 수동 다크모드 클래스 (Gradio 토글용) */
[data-theme="dark"], 
.dark,
.gr-theme-dark {
    /* 배경 컬러 */
    --background-color: #1a1a1a;
    --card-bg: #2d2d2d;
    --input-bg: #2d2d2d;
    
    /* 텍스트 컬러 */
    --text-color: #e5e5e5;
    --text-secondary: #a1a1aa;
    
    /* 보더 및 구분선 */
    --border-color: #404040;
    --border-light: #525252;
    
    /* 테이블 컬러 */
    --table-even-bg: #333333;
    --table-hover-bg: #404040;
    
    /* 그림자 */
    --shadow: 0 8px 30px rgba(0, 0, 0, 0.3);
    --shadow-light: 0 2px 4px rgba(0, 0, 0, 0.2);
}
/* 4. 기본 요소 다크모드 적용 */
body {
    font-family: 'Pretendard', 'Noto Sans KR', -apple-system, BlinkMacSystemFont, sans-serif;
    background-color: var(--background-color) !important;
    color: var(--text-color) !important;
    line-height: 1.6;
    transition: background-color 0.3s ease, color 0.3s ease;
}
/* 5. Gradio 컨테이너 강제 적용 */
.gradio-container,
.gradio-container *,
.gr-app,
.gr-app *,
.gr-interface {
    background-color: var(--background-color) !important;
    color: var(--text-color) !important;
}
/* Gradio 컨테이너 오버라이드 */
.gradio-container {
    max-width: 100% !important;
    width: 100% !important;
    margin: 0 auto !important;
    padding: 0 !important;
    background-color: var(--background-color) !important;
    box-sizing: border-box !important;
}
/* 추가: 내부 컨테이너도 100% 너비로 설정 */
.contain {
    max-width: 100% !important;
    width: 100% !important;
}
/* 추가: 각 행(Row)도 100% 너비로 설정 */
.gr-padded {
    padding: 0 !important;
    width: 100% !important;
    max-width: 100% !important;
}
/* 6. 카드 및 패널 스타일 */
.gr-group,
.gr-form,
.gr-box,
.gr-panel,
.custom-frame,
[class*="frame"],
[class*="card"],
[class*="panel"] {
    background-color: var(--card-bg) !important;
    border-radius: var(--border-radius) !important;
    box-shadow: var(--shadow) !important;
    padding: 1.5rem !important;
    margin-bottom: 1.5rem !important;
    border: 1px solid var(--border-color) !important;
    transition: transform 0.3s ease, background-color 0.3s ease;
    color: var(--text-color) !important;
}
.gr-group:hover {
    transform: translateY(-5px);
}
/* 7. 입력 필드 스타일 */
input[type="text"],
input[type="number"],
input[type="email"],
input[type="password"],
textarea,
select,
.gr-input,
.gr-text-input,
.gr-textarea,
.gr-dropdown {
    background-color: var(--input-bg) !important;
    color: var(--text-color) !important;
    border: 1px solid var(--border-color) !important;
    border-radius: var(--border-radius) !important;
    padding: 12px !important;
    box-shadow: inset 0 1px 3px rgba(0, 0, 0, 0.05) !important;
    transition: all 0.3s ease !important;
}
input[type="text"]:focus,
input[type="number"]:focus,
input[type="email"]:focus,
input[type="password"]:focus,
textarea:focus,
select:focus,
.gr-input:focus,
.gr-text-input:focus,
.gr-textarea:focus,
.gr-dropdown:focus {
    border-color: var(--primary-color) !important;
    outline: none !important;
    box-shadow: 0 0 0 2px rgba(251, 127, 13, 0.2) !important;
}
/* 8. 라벨 및 텍스트 요소 */
label,
.gr-label,
.gr-checkbox label,
.gr-radio label,
p, span, div {
    color: var(--text-color) !important;
}
/* 9. 섹션 제목 */
.section-title {
    font-size: 22px !important;
    font-weight: 700 !important;
    color: var(--text-color) !important;
    margin-bottom: 1rem !important;
    padding-bottom: 0.5rem !important;
    border-bottom: 2px solid var(--primary-color) !important;
    display: flex;
    align-items: center;
}
.section-title span {
    color: var(--primary-color);
}
/* 10. 버튼 스타일링 */
.custom-button {
    background: linear-gradient(135deg, var(--primary-color), var(--secondary-color)) !important;
    color: white !important;
    font-weight: 600 !important;
    border: none !important;
    border-radius: 30px !important;
    padding: 12px 24px !important;
    box-shadow: 0 4px 8px rgba(251, 127, 13, 0.25) !important;
    transition: all 0.3s ease !important;
    text-transform: none !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
}
.custom-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 12px rgba(251, 127, 13, 0.3) !important;
}
.custom-button.primary {
    background: linear-gradient(135deg, var(--accent-color), #ff9a8b) !important;
}
button:not([class*="custom"]):not([class*="primary"]):not([class*="secondary"]) {
    background-color: var(--card-bg) !important;
    color: var(--text-color) !important;
    border: 1px solid var(--border-color) !important;
    border-radius: var(--border-radius) !important;
    transition: all 0.3s ease !important;
}
/* 11. 이미지 컨테이너 */
.image-container {
    border-radius: var(--border-radius);
    overflow: hidden;
    border: 1px solid var(--border-color);
    transition: all 0.3s ease;
    background-color: var(--card-bg);
}
.image-container:hover {
    box-shadow: 0 10px 20px rgba(0, 0, 0, 0.1);
}
/* 12. 테이블 스타일 */
table {
    background-color: var(--card-bg) !important;
    color: var(--text-color) !important;
    border-color: var(--border-color) !important;
}
table th {
    background-color: var(--primary-color) !important;
    color: white !important;
    border-color: var(--border-color) !important;
}
table td {
    background-color: var(--card-bg) !important;
    color: var(--text-color) !important;
    border-color: var(--border-color) !important;
}
table tbody tr:nth-child(even) {
    background-color: var(--table-even-bg) !important;
}
table tbody tr:hover {
    background-color: var(--table-hover-bg) !important;
}
/* 13. 체크박스 및 라디오 버튼 */
input[type="checkbox"],
input[type="radio"] {
    accent-color: var(--primary-color) !important;
}
/* 14. 버튼 그룹 */
.button-grid {
    display: grid;
    grid-template-columns: repeat(auto-fill, minmax(150px, 1fr));
    gap: 0.8rem;
    margin-bottom: 1.2rem;
}
/* 15. 스크롤바 스타일 */
::-webkit-scrollbar-thumb:hover {
    background: var(--secondary-color);
}
/* 16. 아코디언 및 드롭다운 */
details {
    background-color: var(--card-bg) !important;
    border-color: var(--border-color) !important;
    color: var(--text-color) !important;
}
details summary {
    background-color: var(--card-bg) !important;
    color: var(--text-color) !important;
}
/* 17. 툴팁 및 팝업 */
[data-tooltip]:hover::after,
.tooltip,
.popup {
    background-color: var(--card-bg) !important;
    color: var(--text-color) !important;
    border-color: var(--border-color) !important;
    box-shadow: var(--shadow-light) !important;
}
/* 18. 모달 및 오버레이 */
.modal,
.overlay,
[class*="modal"],
[class*="overlay"] {
    background-color: var(--card-bg) !important;
    color: var(--text-color) !important;
    border-color: var(--border-color) !important;
}
/* 19. 추가 Gradio 컴포넌트들 */
.gr-block,
.gr-group,
.gr-row,
.gr-column {
    background-color: var(--background-color) !important;
    color: var(--text-color) !important;
}
/* 20. 코드 블록 및 pre 태그 */
code,
pre,
.code-block {
    background-color: var(--table-even-bg) !important;
    color: var(--text-color) !important;
    border-color: var(--border-color) !important;
}
/* 21. 알림 및 메시지 */
.alert,
.message,
.notification,
[class*="alert"],
[class*="message"],
[class*="notification"] {
    background-color: var(--card-bg) !important;
    color: var(--text-color) !important;
    border-color: var(--border-color) !important;
}
/* 22. 애니메이션 스타일 */
@keyframes fadeIn {
    from { opacity: 0; transform: translateY(10px); }
    to { opacity: 1; transform: translateY(0); }
}
.fade-in {
    animation: fadeIn 0.5s ease-out;
}
/* 23. Examples 섹션 스타일 */
.examples-section {
    display: grid;
    grid-template-columns: repeat(auto-fill, minmax(300px, 1fr));
    gap: 1.5rem;
    margin-top: 1rem;
}
.example-item {
    background-color: var(--card-bg);
    border-radius: var(--border-radius);
    overflow: hidden;
    box-shadow: var(--shadow);
    transition: transform 0.3s ease, background-color 0.3s ease;
}
.example-item:hover {
    transform: translateY(-5px);
}
/* 24. 전환 애니메이션 */
* {
    transition: background-color 0.3s ease, 
                color 0.3s ease, 
                border-color 0.3s ease !important;
}
/* 25. 반응형 */
@media (max-width: 768px) {
    .button-grid {
        grid-template-columns: repeat(2, 1fr);
    }
    .examples-section {
        grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
    }
}
"""

# FontAwesome 아이콘 포함
fontawesome_link = """
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" crossorigin="anonymous" referrerpolicy="no-referrer" />
"""

# 제목과 사용 가이드 제거
header_html = ""
image_generator_guide_html = ""

# UI 구성
with gr.Blocks(css=custom_css, theme=gr.themes.Default(
    primary_hue="orange",
    secondary_hue="orange",
    font=[gr.themes.GoogleFont("Noto Sans KR"), "ui-sans-serif", "system-ui"]
)) as demo:
    gr.HTML(fontawesome_link)
    # 제목 제거
    # gr.HTML(header_html)
    
    # 탭 제거하고 이미지 생성기 내용을 직접 배치
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            # ======== 이미지 업로드 및 설정 섹션 ========
            with gr.Group():
                gr.HTML('<div class="section-title"><i class="fas fa-upload"></i> <span>이미지 업로드 및 설정</span></div>')
                with gr.Row():
                    image1_input = gr.Image(type="pil", label="#1", image_mode="RGB", elem_classes="image-container", height=400)
                    image2_input = gr.Image(type="pil", label="#2", image_mode="RGB", elem_classes="image-container", height=400)
                    image3_input = gr.Image(type="pil", label="#3", image_mode="RGB", elem_classes="image-container", height=400)
                
                # 프롬프트 입력 필드 추가
                prompt_input = gr.Textbox(
                    lines=3,
                    placeholder="프롬프트를 입력하거나 비워두면 자동 합성됩니다. '#1', '#2', '#3'으로 각 이미지를 참조할 수 있습니다.",
                    label="프롬프트 (선택 사항)",
                    elem_classes="gr-text-input"
                )
            
            # ======== 변환 옵션 섹션 ========
            with gr.Group():
                gr.HTML('<div class="section-title"><i class="fas fa-sliders-h"></i> <span>프롬프트 템플릿</span></div>')
                with gr.Column(elem_classes="button-grid"):
                    image_change_btn1 = gr.Button('🔄 부분변경-1', elem_classes="custom-button")
                    image_change_btn2 = gr.Button('🔄 부분변경-2', elem_classes="custom-button")
                    image_change_btn3= gr.Button('🔄 부분변경-3', elem_classes="custom-button")
                    text_remove_btn = gr.Button('🧹 글자지우기', elem_classes="custom-button")
                    text_change_btn = gr.Button('🔤 글자변경', elem_classes="custom-button")
                    clothes_change_btn1 = gr.Button('👕 상품착용-1', elem_classes="custom-button")
                    clothes_change_btn2 = gr.Button('👓 상품착용-2', elem_classes="custom-button")
                    holding_product_btn = gr.Button('🍷 상품들고', elem_classes="custom-button")
                    background_change_btn = gr.Button('🖼️ 배경바꾸기', elem_classes="custom-button")
                    composite_product_btn = gr.Button('✂️ 부분지우기', elem_classes="custom-button")
                    outpainting_btn = gr.Button('🔍 이미지확장', elem_classes="custom-button")
                    food_btn_1 = gr.Button('🍽️ 플레이팅-1', elem_classes="custom-button")
                    food_btn_2 = gr.Button('🍽️ 플레이팅-2', elem_classes="custom-button")
                    food_btn_3 = gr.Button('🍽️ 플레이팅-3', elem_classes="custom-button")
            
            # ======== 이미지 생성 섹션 ========
            with gr.Group():
                gr.HTML('<div class="section-title"><i class="fas fa-image"></i> <span>이미지 생성</span></div>')
                submit_single_btn = gr.Button('✨ 이미지 생성 (1장)', elem_classes="custom-button primary")
                submit_btn = gr.Button('✨ 이미지 생성 (2장)', elem_classes="custom-button primary")
                
        with gr.Column(scale=1):
            # ======== 생성된 이미지 섹션 ========
            with gr.Group():
                gr.HTML('<div class="section-title"><i class="fas fa-images"></i> <span>생성된 이미지</span></div>')
                with gr.Row():
                    with gr.Column():
                        output_image1 = gr.Image(label="이미지 #1", elem_classes="image-container", height=400)
                    with gr.Column():
                        output_image2 = gr.Image(label="이미지 #2", elem_classes="image-container", height=400)
            
            # ======== 결과 정보 섹션 ========
            with gr.Group():
                gr.HTML('<div class="section-title"><i class="fas fa-info-circle"></i> <span>결과 정보</span></div>')
                output_text = gr.Textbox(label="상태 메시지", lines=2, elem_classes="gr-text-input")
                prompt_display = gr.Textbox(label="사용된 프롬프트 (영어)", visible=True, lines=2, elem_classes="gr-text-input")
    
    # ======== 예제 이미지 섹션 ========
    gr.HTML('<div class="section-title"><i class="fas fa-lightbulb"></i> <span>예제 이미지</span></div>')
    
    # 모든 예제 한 페이지에 표시
    examples = [
        ["down/모델.jpg", None, None, "(#1의 여성)이 살짝 뒤로 돌아보는 모습으로 최대한 이전 seed를 유지한체 자연스럽게 변경하라."],
        ["down/상어레고모형.png", None, None, "(#1 레모모형)에서 청색상어레고만 검은색 고래레고로 변경하고 나머지 부분은 seed를 변경을 하지마라."],
        ["down/얼음가방.png", None, None, "(#1 여행용 얼음박스)앞에 얼음이 담긴 3잔의 콜라가 놓여있는 이미지를 생성하라."],
        ["down/중국어.png", None, None, "(#1 이미지)에 있는 중국어를 모두 제거하라."],
        ["down/텍스트.webp", None, None, '(#1의 텍스트)를 스타일을 유지한체 텍스트만 "Hello"로 바꿔라'],
        ["down/모델2.png", "down/선글라스.png", "down/청바지.png", "(#1의 여성모델)이 신체 비율과 포즈는 유지한 체 (#2의 선글라스)와 (#3의 청바지)를 직접 모델이 착용한것 처럼 자연스럽게 모습을 생성하라."],
        ["down/모델2.png", "down/선글라스.png", "down/카페전경.png", "(#1의 여성모델)이 (#2의 선글라스)을 착용하고 (#3의 뒷배경의 카페전체가 보이며) 의자에 앉아 있는 모습을 생성하라."],
        ["down/모델2.png", "down/와인잔.png", None, "(#1의 여성모델)이(#2의 와인잔)을 들고 있는 자연스러운 모습을 생성하라."],
        ["down/모델2.png", "down/카페전경.png", None, "(#1의 여성모델)이 (#2 카페)에서 자연스럽게 있는 모습을 생성하라."],
        ["down/상어레고모형.png", None, None, "(#1의 레고모형)에서 청색상어레고를 제거한 후, 그 자리를 주변 배경과 자연스럽게 어우러지도록 채워주세요. 단, 이미지의 다른 부분의 주요 요소는 동일하게 유지해야한다."],
        ["down/카페전경.png", None, None, "(#1 이미지)를 원본그대로 중앙에 두고 비율로 유지한 체 위아래 및 좌우로 크게 확장하라."],
        ["down/샐러드.png", None, None, "(#1셀러드)에 담은 용기는 버리고 넓고 흰 예쁜 접시에  (#1셀러드)음식만 가득 채워서 상업적인 각도로 어울리는 소품과 함께 플레이팅 한 모습을 이미지로 생성하라. "],
        ["down/샐러드.png", "down/플레이팅.png", None, "(#2 플레이팅한 이미지)에 담긴 음식을  (#1 샐러드)로 바꾸고 나머지는 시드를 유지한 체 이미지를 생성하라."],
        ["down/컵.png", None, None, "(#1컵)에 딸기, 바닐라, 초코 아이스크림을 담고 그 위에 초코 시럽이 흐르게 이미지를 생성하라."]
    ]
    
    # 모든 예제를 한 페이지에 표시하도록 수정된 부분
    gr.Examples(
        examples=examples,
        inputs=[image1_input, image2_input, image3_input, prompt_input],
        examples_per_page=len(examples)  # 모든 예제를 한 페이지에 표시
    )
    
    # ========== 이미지 생성기 이벤트 연결 ==========
    # 버튼 이벤트 연결
    image_change_btn1.click(
        fn=lambda: "(#1의 여성)이 살짝 뒤로 돌아보는 모습으로 최대한 이전 seed를 유지한테 자연스럽게 변경하라.",
        inputs=[],
        outputs=prompt_input
    )
    image_change_btn2.click(
        fn=lambda: "(#1 레모모형)에서 청색상어레고만 검은색 고래레고로 변경하고 나머지 부분은 seed를 변경을 하지마라.",
        inputs=[],
        outputs=prompt_input
    )
    image_change_btn3.click(
        fn=lambda: "(#1 여행용 얼음박스)앞에 얼음이 담긴 3잔의 콜라가 놓여있는 이미지를 생성하라.",
        inputs=[],
        outputs=prompt_input
    )
    
    text_remove_btn.click(
        fn=lambda: "(#1 이미지)에 있는 중국어를 모두 제거하라.",
        inputs=[],
        outputs=prompt_input
    )
    text_change_btn.click(
        fn=lambda: '(#1의 텍스트)를 스타일을 유지한체 텍스트만 "Hello"로 바꿔라',
        inputs=[],
        outputs=prompt_input
    )
    clothes_change_btn1.click(
        fn=lambda: "(#1의 여성모델)이 신체 비율과 포즈는 유지한 체 (#2의 선글라스)와 (#3의 청바지)를 직접 모델이 착용한것 처럼 자연스럽게 모습을 생성하라.",
        inputs=[],
        outputs=prompt_input
    )
    clothes_change_btn2.click(
        fn=lambda: "(#1의 여성모델)이 (#2의 선글라스)을 착용하고 (#3의 뒷배경의 카페전체가 보이며) 의자에 앉아 있는 모습을 생성하라.",
        inputs=[],
        outputs=prompt_input
    )
    holding_product_btn.click(
        fn=lambda: "(#1의 여성모델)이(#2의 와인잔)을 들고 있는 자연스러운 모습을 생성하라.",
        inputs=[],
        outputs=prompt_input
    )
    background_change_btn.click(
        fn=lambda: "(#1의 여성모델)이 (#2 카페)에서 자연스럽게 있는 모습을 생성하라.",
        inputs=[],
        outputs=prompt_input
    )
    composite_product_btn.click(
        fn=lambda: "(#1의 레고모형)에서 청색상어레고를 제거한 후, 그 자리를 주변 배경과 자연스럽게 어우러지도록 채워주세요. 단, 이미지의 다른 부분의 주요 요소는 동일하게 유지 해야한다.",
        inputs=[],
        outputs=prompt_input
    )
    outpainting_btn.click(
        fn=lambda: "(#1 이미지)를 원본그대로 중앙에 두고 비율로 유지한 체 위아래 및 좌우로 크게 확장하라.",
        inputs=[],
        outputs=prompt_input

    )
    food_btn_1.click(
        fn=lambda: "(#1셀러드)에 담은 용기는 버리고 넓고 흰 예쁜 접시에 (#1셀러드)음식만 가득 채워서 상업적인 각도로 어울리는 소품과 함께 플레이팅 한 모습을 이미지로 생성하라. ",
        inputs=[],
        outputs=prompt_input
    )
    food_btn_2.click(
        fn=lambda: "(#2 플레이팅한 이미지)에 담긴 음식을  (#1 샐러드)로 바꾸고 나머지는 시드를 유지한 체 이미지를 생성하라.",
        inputs=[],
        outputs=prompt_input
    )  

    food_btn_3.click(
        fn=lambda: "(#1컵)에 딸기, 바닐라, 초코 아이스크림을 담고 그 위에 초코 시럽이 흐르게 이미지를 생성하라.",
        inputs=[],
        outputs=prompt_input
    )    
    
    # 단일 이미지 생성 버튼 이벤트 연결 - API 키 입력값 제거
    def generate_single_image(image1, image2, image3, prompt):
        return process_images_with_prompt(image1, image2, image3, prompt, 0, 3)
    
    submit_single_btn.click(
        fn=generate_single_image,
        inputs=[image1_input, image2_input, image3_input, prompt_input],
        outputs=[output_image1, output_text, prompt_display],
    )
    
    # 2장 이미지 생성 버튼 이벤트 연결 - API 키 입력값 제거
    submit_btn.click(
        fn=generate_multiple_images,
        inputs=[image1_input, image2_input, image3_input, prompt_input],
        outputs=[output_image1, output_image2, output_text, prompt_display],
    )


# API 키 초기화 및 애플리케이션 실행
initialize_api_keys()  # API 키 초기화 함수 호출
demo.queue()
demo.launch(share=False, inbrowser=True, width="100%")  # width 파라미터 추가