import os
import tempfile
import logging
import re
import time
import json
from PIL import Image
import gradio as gr
from google import genai
from google.genai import types
import google.generativeai as genai_generative
from dotenv import load_dotenv

load_dotenv()

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

# ------------------- 배경 디렉토리 설정 -------------------
BACKGROUNDS_DIR = "./background"
if not os.path.exists(BACKGROUNDS_DIR):
    os.makedirs(BACKGROUNDS_DIR)
    logger.info(f"배경 디렉토리를 생성했습니다: {BACKGROUNDS_DIR}")

# ------------------- 전역 변수 설정 -------------------
SIMPLE_BACKGROUNDS = {}
STUDIO_BACKGROUNDS = {}
NATURE_BACKGROUNDS = {}
INDOOR_BACKGROUNDS = {}
TECHNOLOGY_BACKGROUNDS = {}
COLORFUL_PATTERN_BACKGROUNDS = {}
ABSTRACT_BACKGROUNDS = {}
JEWELRY_BACKGROUNDS = {}  # 쥬얼리 배경 전역 변수 추가

# ------------------- 배경 JSON 파일 로드 함수 -------------------
def load_background_json(filename):
    """배경 JSON 파일 로드 함수"""
    file_path = os.path.join(BACKGROUNDS_DIR, filename)
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            logger.info(f"{filename} 파일을 성공적으로 로드했습니다. {len(data)} 항목 포함.")
            return data
    except FileNotFoundError:
        logger.info(f"{filename} 파일이 없습니다.")
        return {}
    except Exception as e:
        logger.warning(f"{filename} 파일 로드 중 오류 발생: {str(e)}.")
        return {}

# ------------------- 배경 옵션 초기화 함수 -------------------
def initialize_backgrounds():
    """모든 배경 옵션 초기화 함수"""
    global SIMPLE_BACKGROUNDS, STUDIO_BACKGROUNDS, NATURE_BACKGROUNDS, INDOOR_BACKGROUNDS
    global TECHNOLOGY_BACKGROUNDS, COLORFUL_PATTERN_BACKGROUNDS, ABSTRACT_BACKGROUNDS
    global JEWELRY_BACKGROUNDS  # 쥬얼리 배경 추가
    
    logger.info(f"Backgrounds 디렉토리 경로: {BACKGROUNDS_DIR}")
    logger.info(f"디렉토리 내 파일 목록: {os.listdir(BACKGROUNDS_DIR)}")
    
    SIMPLE_BACKGROUNDS = load_background_json("simple_backgrounds.json")
    STUDIO_BACKGROUNDS = load_background_json("studio_backgrounds.json")
    NATURE_BACKGROUNDS = load_background_json("nature_backgrounds.json")
    INDOOR_BACKGROUNDS = load_background_json("indoor_backgrounds.json")
    TECHNOLOGY_BACKGROUNDS = load_background_json("tech-backgrounds-final.json")
    COLORFUL_PATTERN_BACKGROUNDS = load_background_json("colorful-pattern-backgrounds.json")
    ABSTRACT_BACKGROUNDS = load_background_json("abstract_backgrounds.json")
    JEWELRY_BACKGROUNDS = load_background_json("jewelry_backgrounds.json")
    
    # 기본값 설정 (파일이 없거나 비어있는 경우)
    if not SIMPLE_BACKGROUNDS:
        SIMPLE_BACKGROUNDS = {"클래식 화이트": "clean white background with soft even lighting"}
    if not STUDIO_BACKGROUNDS:
        STUDIO_BACKGROUNDS = {"미니멀 플랫레이": "minimalist flat lay with clean white background"}
    if not NATURE_BACKGROUNDS:
        NATURE_BACKGROUNDS = {"열대 해변": "tropical beach with crystal clear water"}
    if not INDOOR_BACKGROUNDS:
        INDOOR_BACKGROUNDS = {"미니멀 스칸디나비안 거실": "minimalist Scandinavian living room"}
    if not TECHNOLOGY_BACKGROUNDS:
        TECHNOLOGY_BACKGROUNDS = {"다이나믹 스플래시": "dynamic water splash interaction with product"}
    if not COLORFUL_PATTERN_BACKGROUNDS:
        COLORFUL_PATTERN_BACKGROUNDS = {"화려한 꽃 패턴": "vibrant floral pattern backdrop"}
    if not ABSTRACT_BACKGROUNDS:
        ABSTRACT_BACKGROUNDS = {"네온 라이트": "neon light abstract background with vibrant glowing elements"}
    if not JEWELRY_BACKGROUNDS:
        JEWELRY_BACKGROUNDS = {"클래식 화이트 실크": "pristine white silk fabric backdrop"}
    
    logger.info("모든 배경 옵션 초기화 완료")

# 배경 드롭다운 초기화를 위한 함수
def initialize_dropdowns():
    """드롭다운 메뉴 초기화 함수"""
    simple_choices = list(SIMPLE_BACKGROUNDS.keys())
    studio_choices = list(STUDIO_BACKGROUNDS.keys())
    nature_choices = list(NATURE_BACKGROUNDS.keys())
    indoor_choices = list(INDOOR_BACKGROUNDS.keys())
    tech_choices = list(TECHNOLOGY_BACKGROUNDS.keys())
    colorful_choices = list(COLORFUL_PATTERN_BACKGROUNDS.keys())
    abstract_choices = list(ABSTRACT_BACKGROUNDS.keys())
    jewelry_choices = list(JEWELRY_BACKGROUNDS.keys())
    
    return {
        "simple": simple_choices,
        "studio": studio_choices,
        "nature": nature_choices,
        "indoor": indoor_choices,
        "tech": tech_choices,
        "colorful": colorful_choices,
        "abstract": abstract_choices,
        "jewelry": jewelry_choices,
    }

# ------------------- 기본 유틸리티 함수 -------------------
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")
    try:
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            logger.error("Gemini API 키가 설정되지 않았습니다.")
            prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
            return prompt
        client = genai.Client(api_key=api_key)
        translation_prompt = f"""
        Translate the following Korean text to English:
        
        {prompt}
        
        IMPORTANT: The token IMAGE_TAG_ONE is a special tag 
        and must be preserved exactly as is in your translation. Do not translate this token.
        """
        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")
            logger.info(f"Translated text: {translated_text.strip()}")
            return translated_text.strip()
        else:
            logger.warning("번역 결과가 없습니다. 원본 프롬프트 사용")
            prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
            return prompt
    except Exception as e:
        logger.exception("번역 중 오류 발생:")
        prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
        return prompt

def preprocess_prompt(prompt, image1):
    has_img1 = image1 is not None
    if "#1" in prompt and not has_img1:
        prompt = prompt.replace("#1", "첫 번째 이미지(없음)")
    else:
        prompt = prompt.replace("#1", "첫 번째 이미지")
    prompt += " 이미지를 생성해주세요. 이미지에 텍스트나 글자를 포함하지 마세요."
    return prompt

# ------------------- 이미지 생성 함수 -------------------
def generate_with_images(prompt, images, variation_index=0):
    try:
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."
        client = genai.Client(api_key=api_key)
        logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}, 변형 인덱스: {variation_index}")
        
        variation_suffixes = [
            " Create this as a professional studio product shot with precise focus on the product details. Do not add any text, watermarks, or labels to the image.",
            " Create this as a high-contrast artistic studio shot with dramatic lighting and shadows. Do not add any text, watermarks, or labels to the image.",
            " Create this as a soft-lit elegantly styled product shot with complementary elements. Do not add any text, watermarks, or labels to the image.",
            " Create this as a high-definition product photography with perfect color accuracy and detail preservation. 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 + " Create as high-end commercial product photography. 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.05,
                top_p=0.97,
                top_k=50,
                max_output_tokens=10240
            )
        )
        
        with tempfile.NamedTemporaryFile(suffix=".jpg", 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")
            result_img.save(temp_path, format="JPEG", quality=95)
            return temp_path, f"이미지가 성공적으로 생성되었습니다. {result_text}"
    except Exception as e:
        logger.exception("이미지 생성 중 오류 발생:")
        return None, f"오류 발생: {str(e)}"

def process_images_with_prompt(image1, prompt, variation_index=0, max_retries=3):
    retry_count = 0
    last_error = None
    while retry_count < max_retries:
        try:
            images = [image1]
            valid_images = [img for img in images if img is not None]
            if not valid_images:
                return None, "이미지를 업로드해주세요.", ""
            final_prompt = prompt.strip()
            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 filter_prompt_only(prompt):
    """Gemini의 설명 및 불필요한 메시지를 제거하고 실제 프롬프트만 추출하는 함수"""
    code_block_pattern = r"```\s*(.*?)```"
    code_match = re.search(code_block_pattern, prompt, re.DOTALL)
    if code_match:
        return code_match.group(1).strip()
    
    if "--ar 1:1" in prompt:
        lines = prompt.split('\n')
        prompt_lines = []
        in_prompt = False
        for line in lines:
            if (not in_prompt and 
                ("product" in line.lower() or 
                 "magazine" in line.lower() or 
                 "commercial" in line.lower() or 
                 "photography" in line.lower())):
                in_prompt = True
                prompt_lines.append(line)
            elif in_prompt:
                if "explanation" in line.lower() or "let me know" in line.lower():
                    break
                prompt_lines.append(line)
        if prompt_lines:
            return '\n'.join(prompt_lines).strip()
    
    return prompt.strip()

def get_selected_background_info(bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry):
    """선택된 배경 정보를 가져오는 함수"""
    if bg_type == "심플 배경":
        return {
            "category": "심플 배경",
            "name": simple,
            "english": SIMPLE_BACKGROUNDS.get(simple, "white background")
        }
    elif bg_type == "스튜디오 배경":
        return {
            "category": "스튜디오 배경",
            "name": studio,
            "english": STUDIO_BACKGROUNDS.get(studio, "product photography studio")
        }
    elif bg_type == "자연 환경":
        return {
            "category": "자연 환경",
            "name": nature,
            "english": NATURE_BACKGROUNDS.get(nature, "natural environment")
        }
    elif bg_type == "실내 환경":
        return {
            "category": "실내 환경",
            "name": indoor,
            "english": INDOOR_BACKGROUNDS.get(indoor, "indoor environment")
        }
    elif bg_type == "테크놀로지 배경":
        return {
            "category": "테크놀로지 배경",
            "name": tech,
            "english": TECHNOLOGY_BACKGROUNDS.get(tech, "technology environment")
        }
    elif bg_type == "컬러풀 패턴 배경":
        return {
            "category": "컬러풀 패턴 배경",
            "name": colorful,
            "english": COLORFUL_PATTERN_BACKGROUNDS.get(colorful, "colorful pattern background")
        }
    elif bg_type == "추상/특수 배경":
        return {
            "category": "추상/특수 배경",
            "name": abstract,
            "english": ABSTRACT_BACKGROUNDS.get(abstract, "abstract background")
        }
    elif bg_type == "쥬얼리 배경":
        return {
            "category": "쥬얼리 배경",
            "name": jewelry,
            "english": JEWELRY_BACKGROUNDS.get(jewelry, "jewelry backdrop")
        }
    else:
        return {
            "category": "기본 배경",
            "name": "화이트 배경",
            "english": "white background"
        }

def generate_enhanced_system_instruction():
    """향상된 시스템 인스트럭션 생성 함수"""
    return """당신은 상품 이미지의 배경을 변경하기 위한 최고 품질의 프롬프트를 생성하는 전문가입니다.
사용자가 제공하는 상품명, 배경 유형, 추가 요청사항을 바탕으로 미드저니(Midjourney)에 사용할 수 있는 상세하고 전문적인 프롬프트를 영어로 생성해주세요.
다음 가이드라인을 반드시 준수해야 합니다:
1. 상품을 "#1"로 지정하여 참조합니다. (예: "skincare tube (#1)")
2. *** 매우 중요: 상품의 원래 특성(디자인, 색상, 형태, 로고, 패키지 등)은 어떤 상황에서도 절대 변경하지 않습니다. ***
3. *** 상품의 본질적 특성을 유지하되, 상품에 포커스를 맞춰 모든 세부 사항이 선명하게 드러나도록 하며, 
   8K 해상도(8K resolution), 오버샤프닝 없는 초고화질(ultra high definition without oversharpening)로 렌더링되어야 합니다. ***
4. 이미지 비율은 정확히 1:1(정사각형) 형식으로 지정합니다. 프롬프트에 "square format", "1:1 ratio" 또는 "aspect ratio 1:1"을 명시적으로 포함합니다.
5. 상품은 반드시 정사각형 구도의 정중앙에 배치되어야 하며, 적절한 크기로 표현하여 디테일이 완벽하게 보이도록 합니다.
6. 상품을 이미지의 주요 초점으로 부각시키고, 상품의 비율이 전체 이미지에서 60-70% 이상 차지하도록 합니다.
7. 조명 설명을 매우 구체적으로 해주세요. 예: "soft directional lighting from left side", "dramatic rim lighting", "diffused natural light through windows"
8. 배경의 재질과 질감을 상세히 설명해주세요. 예: "polished marble surface", "rustic wooden table with visible grain", "matte concrete wall with subtle texture"
9. 프롬프트에 다음 요소들을 명시적으로 포함하되, 사용 맥락에 적절하게 변형하세요:
   - "award-winning product photography" 
   - "magazine-worthy commercial product shot"
   - "professional advertising imagery with perfect exposure"
   - "studio lighting with color-accurate rendering"
   - "8K ultra high definition product showcase"
   - "commercial product photography with precise detail rendering"
   - "ultra high definition"
   - "crystal clear details"
10. 사용자가 제공한 구체적인 배경과 추가 요청사항을 프롬프트에 정확히 반영하고 확장합니다.
11. 프롬프트 끝에 미드저니 파라미터 "--ar 1:1 --s 750 --q 2 --v 5.2" 파라미터를 추가하여 미드저니에서 고품질 정사각형 비율을 강제합니다.
12. 매우 중요: 프롬프트 외에 다른 설명이나 메타 텍스트를 포함하지 마세요. 오직 프롬프트 자체만 제공하세요.
"""

def generate_prompt_with_gemini(product_name, background_info, additional_info=""):
    """향상된 프롬프트 생성 함수"""
    GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
    if not GEMINI_API_KEY:
        return "Gemini API 키가 설정되지 않았습니다. 환경 변수 GEMINI_API_KEY를 설정하거나 코드에 직접 입력하세요."
    
    try:
        genai_generative.configure(api_key=GEMINI_API_KEY)
        
        prompt_request = f"""
상품명: {product_name}
배경 유형: {background_info.get('english', 'studio')}
배경 카테고리: {background_info.get('category', '')}
배경 이름: {background_info.get('name', '')}
추가 요청사항: {additional_info}
중요 요구사항:
1. 상품(#1)이 이미지 구도에서 중심적인 위치를 차지하며 적절한 크기(이미지의 60-70%)로 표현되도록 프롬프트를 생성해주세요.
2. 이미지는 정확히 1:1 비율(정사각형)이어야 합니다.
3. 상품의 디자인, 색상, 형태, 로고 등 본질적 특성은 절대 수정하지 마세요.
4. 구체적인 조명 기법을 상세히 명시해주세요:
   - 정확한 조명 위치 (예: "45-degree key light from upper left")
   - 조명 품질 (예: "soft diffused light", "hard directional light")
   - 조명 강도와 색온도 (예: "warm tungsten key light with cool blue fill")
   - 반사와 그림자 처리 방식 (예: "controlled specular highlights with soft shadow transitions")
5. 상품을 더 돋보이게 하는 보조 요소(props)를 자연스럽게 활용하되, 상품이 항상 주인공이어야 합니다.
6. 배경 재질과 표면 질감을 구체적으로 설명하고, 상품과의 상호작용 방식을 명시해주세요.
7. 색상 구성(color palette, color harmonies)을 명확히 해주세요.
8. 고급스러운 상업 광고 품질의 이미지가 되도록 프롬프트를 작성해주세요.
9. 프롬프트 끝에 미드저니 파라미터 "--ar 1:1 --s 750 --q 2 --v 5.2"를 추가해주세요.
한국어 입력 내용을 전문적인 영어로 번역하여 반영해주세요.
"""
        model = genai_generative.GenerativeModel(
            'gemini-2.0-flash',
            system_instruction=generate_enhanced_system_instruction()
        )
        
        response = model.generate_content(
            prompt_request,
            generation_config=genai_generative.types.GenerationConfig(
                temperature=0.8,
                top_p=0.97,
                top_k=64,
                max_output_tokens=1600,
            )
        )
        
        response_text = response.text.strip()
        
        if "--ar 1:1" not in response_text:
            response_text = response_text.rstrip(".") + ". --ar 1:1 --s 750 --q 2 --v 5.2"
            
        return response_text
    except Exception as e:
        return f"프롬프트 생성 중 오류가 발생했습니다: {str(e)}"

# ------------------- 단일 이미지 생성 함수 -------------------
def generate_product_image(image, bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry, product_name, additional_info):
    if image is None:
        return None, "이미지를 업로드해주세요.", "이미지를 업로드 후 프롬프트를 생성해주세요."
    product_name = product_name.strip() or "제품"
    background_info = get_selected_background_info(bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry)
    generated_prompt = generate_prompt_with_gemini(product_name, background_info, additional_info)
    if "Gemini API 키가 설정되지 않았습니다" in generated_prompt:
        warning_msg = (
            "[Gemini API 키 누락]\n"
            "API 키 설정 방법:\n"
            "1. 환경 변수: export GEMINI_API_KEY=\"your-api-key\"\n"
            "2. 코드 내 직접 입력: GEMINI_API_KEY = \"your-api-key\"\n"
            "키 발급: https://aistudio.google.com/apikey"
        )
        return None, warning_msg, warning_msg
    final_prompt = filter_prompt_only(generated_prompt)
    result_image, status, _ = process_images_with_prompt(image, final_prompt, 0)
    return result_image, status, final_prompt

# ------------------- 4장 이미지 생성 함수 -------------------
def generate_product_images(image, bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry, product_name, additional_info):
    if image is None:
        return None, None, None, None, "이미지를 업로드해주세요.", "이미지를 업로드 후 프롬프트를 생성해주세요."
    product_name = product_name.strip() or "제품"
    background_info = get_selected_background_info(bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry)
    generated_prompt = generate_prompt_with_gemini(product_name, background_info, additional_info)
    if "Gemini API 키가 설정되지 않았습니다" in generated_prompt:
        warning_msg = (
            "[Gemini API 키 누락]\n"
            "API 키 설정 방법:\n"
            "1. 환경 변수: export GEMINI_API_KEY=\"your-api-key\"\n"
            "2. 코드 내 직접 입력: GEMINI_API_KEY = \"your-api-key\"\n"
            "키 발급: https://aistudio.google.com/apikey"
        )
        return None, None, None, None, warning_msg, warning_msg
    final_prompt = filter_prompt_only(generated_prompt)
    images_list = []
    statuses = []
    for i in range(4):
        result_img, status, _ = process_images_with_prompt(image, final_prompt, variation_index=i)
        images_list.append(result_img)
        statuses.append(f"이미지 #{i+1}: {status}")
        time.sleep(1)
    combined_status = "\n".join(statuses)
    return images_list[0], images_list[1], images_list[2], images_list[3], combined_status, final_prompt

# ------------------- Gradio 인터페이스 구성 -------------------
def create_app():
    dropdown_options = initialize_dropdowns()
    
    with gr.Blocks(title="고급 상품 이미지 배경 프롬프트 및 이미지 생성") as demo:
        gr.Markdown("# 고급 상품 이미지 배경 프롬프트 및 이미지 생성")
        gr.Markdown(
            "상품 이미지를 업로드하고, 상품명, 배경 옵션, 추가 요청사항 및 Gemini API 키를 입력하면 Gemini API를 통해 영어 프롬프트를 생성하고, 해당 프롬프트로 이미지가 생성됩니다.\n\n"
            "단일 이미지 생성과 4장 이미지 생성 모두 가능합니다.\n\n"
            "[Gemini API 키 받기](https://aistudio.google.com/apikey)"
        )
        
        with gr.Row():
            # API 키 입력 필드만 표시
            gemini_api_key = gr.Textbox(
                label="Gemini API 키", 
                type="password", 
                placeholder="API 키를 입력하세요",
                interactive=True
            )
        
        with gr.Row():
            with gr.Column(scale=1):
                product_name = gr.Textbox(label="상품명 (한국어 입력)", placeholder="예: 스킨케어 튜브, 스마트워치, 향수, 운동화 등", interactive=True)
                image_input = gr.Image(label="상품 이미지 업로드", type="pil")
                background_type = gr.Radio(
                    choices=["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "테크놀로지 배경", "컬러풀 패턴 배경", "추상/특수 배경", "쥬얼리 배경"],
                    label="배경 유형",
                    value="심플 배경"
                )
                simple_dropdown = gr.Dropdown(
                    choices=dropdown_options["simple"],
                    value=dropdown_options["simple"][0] if dropdown_options["simple"] else None,
                    label="심플 배경 선택",
                    visible=True,
                    interactive=True
                )
                studio_dropdown = gr.Dropdown(
                    choices=dropdown_options["studio"],
                    value=dropdown_options["studio"][0] if dropdown_options["studio"] else None,
                    label="스튜디오 배경 선택",
                    visible=False,
                    interactive=True
                )
                nature_dropdown = gr.Dropdown(
                    choices=dropdown_options["nature"],
                    value=dropdown_options["nature"][0] if dropdown_options["nature"] else None,
                    label="자연 환경 선택",
                    visible=False,
                    interactive=True
                )
                indoor_dropdown = gr.Dropdown(
                    choices=dropdown_options["indoor"],
                    value=dropdown_options["indoor"][0] if dropdown_options["indoor"] else None,
                    label="실내 환경 선택",
                    visible=False,
                    interactive=True
                )
                tech_dropdown = gr.Dropdown(
                    choices=dropdown_options["tech"],
                    value=dropdown_options["tech"][0] if dropdown_options["tech"] else None,
                    label="테크놀로지 배경 선택",
                    visible=False,
                    interactive=True
                )
                colorful_dropdown = gr.Dropdown(
                    choices=dropdown_options["colorful"],
                    value=dropdown_options["colorful"][0] if dropdown_options["colorful"] else None,
                    label="컬러풀 패턴 배경 선택",
                    visible=False,
                    interactive=True
                )
                abstract_dropdown = gr.Dropdown(
                    choices=dropdown_options["abstract"],
                    value=dropdown_options["abstract"][0] if dropdown_options["abstract"] else None,
                    label="추상/특수 배경 선택",
                    visible=False,
                    interactive=True
                )
                jewelry_dropdown = gr.Dropdown(
                    choices=dropdown_options["jewelry"],
                    value=dropdown_options["jewelry"][0] if dropdown_options["jewelry"] else None,
                    label="쥬얼리 배경 선택",
                    visible=False,
                    interactive=True
                )
                additional_info = gr.Textbox(
                    label="추가 요청사항 (선택사항)",
                    placeholder="예: 고급스러운 느낌, 밝은 조명, 자연스러운 보조 객체 등",
                    lines=3,
                    interactive=True
                )
                
                def update_dropdowns(bg_type):
                    return {
                        simple_dropdown: gr.update(visible=(bg_type == "심플 배경")),
                        studio_dropdown: gr.update(visible=(bg_type == "스튜디오 배경")),
                        nature_dropdown: gr.update(visible=(bg_type == "자연 환경")),
                        indoor_dropdown: gr.update(visible=(bg_type == "실내 환경")),
                        tech_dropdown: gr.update(visible=(bg_type == "테크놀로지 배경")),
                        colorful_dropdown: gr.update(visible=(bg_type == "컬러풀 패턴 배경")),
                        abstract_dropdown: gr.update(visible=(bg_type == "추상/특수 배경")),
                        jewelry_dropdown: gr.update(visible=(bg_type == "쥬얼리 배경"))
                    }
                background_type.change(
                    fn=update_dropdowns,
                    inputs=[background_type],
                    outputs=[simple_dropdown, studio_dropdown, nature_dropdown, indoor_dropdown, tech_dropdown, colorful_dropdown, abstract_dropdown, jewelry_dropdown]
                )
                
                with gr.Row():
                    single_btn = gr.Button("프롬프트 및 단일 이미지 생성", variant="primary")
                    multi_btn = gr.Button("프롬프트 및 4장 이미지 생성", variant="primary")
            
            with gr.Column(scale=1):
                # 프롬프트 출력 UI 제거, 이미지 4장만 출력
                with gr.Row():
                    image_output1 = gr.Image(label="이미지 #1", type="filepath")
                    image_output2 = gr.Image(label="이미지 #2", type="filepath")
                with gr.Row():
                    image_output3 = gr.Image(label="이미지 #3", type="filepath")
                    image_output4 = gr.Image(label="이미지 #4", type="filepath")
                status_output = gr.Textbox(label="결과 정보", lines=3)
        
        # 단일 이미지 생성을 위한 함수 (API 키 입력값 활용)
        def modified_single_image_gen(api_key, image, bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry, product_name, additional_info):
            os.environ["GEMINI_API_KEY"] = api_key.strip()
            result_img, status, _ = generate_product_image(image, bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry, product_name, additional_info)
            return result_img, None, None, None, status
        
        single_btn.click(
            fn=modified_single_image_gen,
            inputs=[gemini_api_key, image_input, background_type, simple_dropdown, studio_dropdown, nature_dropdown, indoor_dropdown, tech_dropdown, colorful_dropdown, abstract_dropdown, jewelry_dropdown, product_name, additional_info],
            outputs=[image_output1, image_output2, image_output3, image_output4, status_output]
        )
        
        # 4장 이미지 생성을 위한 함수 (API 키 입력값 활용)
        def modified_multi_image_gen(api_key, image, bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry, product_name, additional_info):
            os.environ["GEMINI_API_KEY"] = api_key.strip()
            img1, img2, img3, img4, status, _ = generate_product_images(image, bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry, product_name, additional_info)
            return img1, img2, img3, img4, status
        
        multi_btn.click(
            fn=modified_multi_image_gen,
            inputs=[gemini_api_key, image_input, background_type, simple_dropdown, studio_dropdown, nature_dropdown, indoor_dropdown, tech_dropdown, colorful_dropdown, abstract_dropdown, jewelry_dropdown, product_name, additional_info],
            outputs=[image_output1, image_output2, image_output3, image_output4, status_output]
        )
    
    return demo

# ------------------- 메인 실행 함수 -------------------
if __name__ == "__main__":
    initialize_backgrounds()
    app = create_app()
    app.queue()
    app.launch()