import os
import tempfile
from PIL import Image
import gradio as gr
import logging
import re
import time

from google import genai
from google.genai import types
from dotenv import load_dotenv
from transformers import pipeline  # 새로 추가된 import

load_dotenv()

# 기존 코드 유지 (로깅, 함수 등 모든 기능 코드)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 배경 제거 기능 추가
def remove_background(image):
    if image is None:
        return None, "이미지가 업로드되지 않았습니다."
    
    try:
        logger.info("배경 제거 시작")
        
        # 모델을 처음 로드할 때는 다운로드할 수 있으므로 시간이 걸릴 수 있음
        pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
        
        # 배경이 제거된 이미지 가져오기
        output_image = pipe(image)
        
        logger.info("배경 제거 완료")
        return output_image, "배경이 성공적으로 제거되었습니다."
    except Exception as e:
        logger.exception("배경 제거 중 오류 발생:")
        return None, f"오류 발생: {str(e)}"


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 = os.environ.get("GEMINI_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 = 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 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.",
            " Create this as the third variation with a more creative style. Do not add any text, watermarks, or labels to the image.",
            " Create this as the fourth variation with enhanced details. 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:
                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 = 4
    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) < 4:
        results.append(None)
    
    combined_status = "\n".join(statuses)
    combined_prompts = "\n".join(prompts)
    
    return results[0], results[1], results[2], results[3], combined_status, combined_prompts

# GFPGAN 관련 코드 유지
import sys
from torchvision.transforms import functional
sys.modules["torchvision.transforms.functional_tensor"] = functional

from basicsr.archs.srvgg_arch import SRVGGNetCompact
from gfpgan.utils import GFPGANer
from realesrgan.utils import RealESRGANer

import torch
import cv2

# 필요한 모델 다운로드
if not os.path.exists('realesr-general-x4v3.pth'):
    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
if not os.path.exists('GFPGANv1.4.pth'):
    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
if not os.path.exists('RestoreFormer.pth'):
    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth -P .")

model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
model_path = 'realesr-general-x4v3.pth'
half = True if torch.cuda.is_available() else False
upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)

def upscaler(img, version, scale):
    try:
        img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
        if len(img.shape) == 3 and img.shape[2] == 4:
            img_mode = 'RGBA'
        elif len(img.shape) == 2:
            img_mode = None
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        else:
            img_mode = None

        h, w = img.shape[0:2]
        if h < 300:
            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)

        face_enhancer = GFPGANer(
            model_path=f'{version}.pth',
            upscale=2,
            arch='RestoreFormer' if version=='RestoreFormer' else 'clean',
            channel_multiplier=2,
            bg_upsampler=upsampler
        )

        try:
            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
        except RuntimeError as error:
            print('오류 발생:', error)

        try:
            if scale != 2:
                interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
                h, w = img.shape[0:2]
                output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
        except Exception as error:
            print('잘못된 스케일 입력:', error)

        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
        return output
    except Exception as error:
        print('전체 예외 발생:', error)
        return None

def upscaler_korean(img):
    return upscaler(img, "GFPGANv1.4", 2)

# 커스텀 CSS 스타일
custom_css = """
:root {
    --primary-color: #5561e9;
    --secondary-color: #6c8aff;
    --accent-color: #ff6b6b;
    --background-color: #f0f5ff;
    --card-bg: #ffffff;
    --text-color: #334155;
    --border-radius: 18px;
    --shadow: 0 8px 30px rgba(0, 0, 0, 0.08);
}

body {
    font-family: 'Pretendard', 'Noto Sans KR', -apple-system, BlinkMacSystemFont, sans-serif;
    background-color: var(--background-color);
    color: var(--text-color);
    line-height: 1.6;
}

/* Gradio 컨테이너 오버라이드 */
.gradio-container {
    max-width: 100% !important;
    margin: 0 auto !important;
    padding: 0 !important;
    background-color: var(--background-color) !important;
}

/* 상단 헤더 */
.app-header {
    background: linear-gradient(135deg, var(--primary-color), var(--secondary-color));
    color: white;
    padding: 2rem;
    border-radius: var(--border-radius);
    margin-bottom: 1.5rem;
    box-shadow: var(--shadow);
    text-align: center;
}

.app-header h1 {
    margin: 0;
    font-size: 2.5rem;
    font-weight: 700;
    letter-spacing: -0.5px;
}

.app-header p {
    margin: 0.75rem 0 0;
    font-size: 1.1rem;
    opacity: 0.9;
}

/* 패널 스타일링 */
.panel {
    background-color: var(--card-bg);
    border-radius: var(--border-radius);
    box-shadow: var(--shadow);
    padding: 1.5rem;
    margin-bottom: 1.5rem;
    border: 1px solid rgba(0, 0, 0, 0.04);
    transition: transform 0.3s ease;
}

.panel:hover {
    transform: translateY(-5px);
}

/* 섹션 제목 */
.section-title {
    font-size: 1.5rem;
    font-weight: 700;
    color: var(--primary-color);
    margin-bottom: 1rem;
    padding-bottom: 0.5rem;
    border-bottom: 2px solid var(--secondary-color);
    display: flex;
    align-items: center;
}

.section-title i {
    margin-right: 0.5rem;
    font-size: 1.4rem;
}

/* 버튼 스타일링 */
.custom-button {
    background: linear-gradient(135deg, var(--primary-color), var(--secondary-color));
    color: white !important;
    border: none !important;
    border-radius: calc(var(--border-radius) - 5px) !important;
    padding: 0.8rem 1.2rem !important;
    font-weight: 600 !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 6px rgba(50, 50, 93, 0.11), 0 1px 3px rgba(0, 0, 0, 0.08);
    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 7px 14px rgba(50, 50, 93, 0.1), 0 3px 6px rgba(0, 0, 0, 0.08) !important;
}

.custom-button.primary {
    background: linear-gradient(135deg, var(--accent-color), #ff9a8b) !important;
}

.custom-button i {
    margin-right: 0.5rem;
}

/* 이미지 컨테이너 */
.image-container {
    border-radius: var(--border-radius);
    overflow: hidden;
    border: 1px solid rgba(0, 0, 0, 0.08);
    transition: all 0.3s ease;
    background-color: white;
}

.image-container:hover {
    box-shadow: 0 10px 20px rgba(0, 0, 0, 0.1);
}

/* 탭 컨테이너 */
.custom-tabs {
    background-color: transparent !important;
    border: none !important;
    margin-bottom: 1rem;
}

.custom-tabs button {
    background-color: rgba(255, 255, 255, 0.7) !important;
    border: none !important;
    border-radius: var(--border-radius) var(--border-radius) 0 0 !important;
    padding: 0.8rem 1.5rem !important;
    font-weight: 600 !important;
    color: var(--text-color) !important;
    transition: all 0.3s ease !important;
}

.custom-tabs button[aria-selected="true"] {
    background-color: var(--primary-color) !important;
    color: white !important;
}

/* 입력 필드 */
.custom-input {
    border-radius: calc(var(--border-radius) - 5px) !important;
    border: 1px solid rgba(0, 0, 0, 0.1) !important;
    padding: 0.8rem 1rem !important;
    transition: all 0.3s ease !important;
}

.custom-input:focus {
    border-color: var(--primary-color) !important;
    box-shadow: 0 0 0 2px rgba(85, 97, 233, 0.2) !important;
}

/* 사용자 매뉴얼 */
.user-manual {
    background-color: white;
    padding: 2rem;
    border-radius: var(--border-radius);
    box-shadow: var(--shadow);
    margin-top: 2rem;
}

.manual-title {
    font-size: 1.8rem;
    font-weight: 700;
    color: var(--primary-color);
    margin-bottom: 1.5rem;
    text-align: center;
    display: flex;
    align-items: center;
    justify-content: center;
}

.manual-title i {
    margin-right: 0.5rem;
    font-size: 1.8rem;
}

.manual-section {
    margin-bottom: 1.5rem;
    padding: 1.2rem;
    background-color: #f8faff;
    border-radius: calc(var(--border-radius) - 5px);
}

.manual-section-title {
    font-size: 1.3rem;
    font-weight: 700;
    margin-bottom: 1rem;
    color: var(--primary-color);
    display: flex;
    align-items: center;
}

.manual-section-title i {
    margin-right: 0.5rem;
    font-size: 1.2rem;
}

.manual-text {
    font-size: 1rem;
    line-height: 1.7;
}

.manual-text strong {
    color: var(--accent-color);
}

.tip-box {
    background-color: rgba(255, 107, 107, 0.1);
    border-left: 3px solid var(--accent-color);
    padding: 1rem 1.2rem;
    margin: 1rem 0;
    border-radius: 8px;
}

/* 버튼 그룹 */
.button-grid {
    display: grid;
    grid-template-columns: repeat(auto-fill, minmax(150px, 1fr));
    gap: 0.8rem;
    margin-bottom: 1.2rem;
}

/* 로딩 애니메이션 */
.progress-container {
    background-color: rgba(255, 255, 255, 0.9);
    border-radius: var(--border-radius);
    padding: 2rem;
    box-shadow: var(--shadow);
    text-align: center;
}

.progress-bar {
    height: 8px;
    border-radius: 4px;
    background: linear-gradient(90deg, var(--primary-color), var(--secondary-color));
    margin: 1rem 0;
}

/* 예시 그리드 */
.examples-grid {
    display: grid;
    grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
    gap: 1rem;
    margin: 1.5rem 0;
}

.example-card {
    background: white;
    border-radius: var(--border-radius);
    overflow: hidden;
    box-shadow: var(--shadow);
    transition: transform 0.3s ease;
}

.example-card:hover {
    transform: translateY(-5px);
}

/* 반응형 */
@media (max-width: 768px) {
    .button-grid {
        grid-template-columns: repeat(2, 1fr);
    }
}

/* 전체 인터페이스 둥근 모서리 */
.block, .prose, .gr-prose, .gr-form, .gr-panel {
    border-radius: var(--border-radius) !important;
}

/* 메인 컨텐츠 스크롤바 */
::-webkit-scrollbar {
    width: 8px;
    height: 8px;
}

::-webkit-scrollbar-track {
    background: rgba(0, 0, 0, 0.05);
    border-radius: 10px;
}

::-webkit-scrollbar-thumb {
    background: var(--secondary-color);
    border-radius: 10px;
}
"""

# 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" />
"""

# 앱 헤더 HTML
header_html = """
<div class="app-header">
    <h1>✨ 이커머스 전용 이미지 생성기(Ver1.1) ✨</h1>
    <p>쉽고 빠르게 상품 이미지를 편집하고 생성하세요! 한 번의 클릭으로 전문적인 이미지를 만들 수 있습니다.</p>
</div>
"""

# 사용자 매뉴얼 HTML
user_manual_html = """
<div class="user-manual">
    <div class="manual-title"><i class="fas fa-book"></i> 사용 설명서</div>
    
    <div class="manual-section">
        <div class="manual-section-title"><i class="fas fa-magic"></i> 이커머스 전용 이미지 생성기</div>
        <p class="manual-text">
            <span style="font-size: 1.1rem; font-weight: 600; color: #5561e9;">1️⃣ 이미지 업로드</span><br>
            최대 3개의 이미지를 업로드할 수 있습니다. 첫 번째 이미지는 필수입니다.<br><br>
            
            <span style="font-size: 1.1rem; font-weight: 600; color: #5561e9;">2️⃣ 프롬프트 입력 또는 버튼 선택</span><br>
            변환하고 싶은 내용을 직접 입력하거나, 미리 정의된 버튼 중 하나를 선택하세요.<br><br>
            
            <span style="font-size: 1.1rem; font-weight: 600; color: #5561e9;">3️⃣ 이미지 생성</span><br>
            '이미지 생성' 버튼을 클릭하면 4가지 버전의 이미지가 생성됩니다.
        </p>
    </div>
    
    <div class="manual-section">
        <div class="manual-section-title"><i class="fas fa-hashtag"></i> 이미지 참조 방법</div>
        <p class="manual-text">
            프롬프트에서 다음 태그를 사용하여 각 이미지를 참조할 수 있습니다:
            <ul>
                <li><strong style="font-size: 1.05rem;">#1</strong> - 첫 번째 이미지</li>
                <li><strong style="font-size: 1.05rem;">#2</strong> - 두 번째 이미지</li>
                <li><strong style="font-size: 1.05rem;">#3</strong> - 세 번째 이미지</li>
            </ul>
            
            <div class="tip-box">
                <i class="fas fa-lightbulb"></i> <strong>팁:</strong> 예를 들어, "(#1의 여성모델)이 (#2의 선글라스)를 착용한 모습"과 같이 사용할 수 있습니다.
            </div>
        </p>
    </div>
    
    <div class="manual-section">
        <div class="manual-section-title"><i class="fas fa-tools"></i> 주요 기능</div>
        <div class="manual-text">
            <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)); gap: 1rem;">
                <div>
                    <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-exchange-alt"></i> 이미지 변경</p>
                    포즈나 스타일을 변경하면서 원본의 주요 요소는 유지합니다.
                </div>
                
                <div>
                    <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-eraser"></i> 글자 지우기/변경</p>
                    이미지의 텍스트를 자연스럽게 제거하거나 다른 텍스트로 변경합니다.
                </div>
                
                <div>
                    <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-tshirt"></i> 가상 상품착용</p>
                    모델에게 다른 이미지의 상품을 자연스럽게 착용시킵니다.
                </div>
                
                <div>
                    <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-image"></i> 배경 바꾸기</p>
                    주요 피사체는 유지하며 배경만 교체합니다.
                </div>
            </div>
        </div>
    </div>
    
    <div class="manual-section">
        <div class="manual-section-title"><i class="fas fa-search-plus"></i> 이미지 업스케일러</div>
        <p class="manual-text">
            <span style="font-size: 1.1rem; font-weight: 600; color: #5561e9;">이미지 개선 기능</span><br>
            저해상도 이미지를 고해상도로 변환하고, 얼굴을 자연스럽게 복원합니다.<br>
            특히 다음과 같은 경우에 효과적입니다:
            <ul>
                <li>해상도가 낮은 이미지 개선</li>
                <li>흐릿하거나 노이즈가 있는 얼굴 복원</li>
                <li>이미지 전반적인 선명도 향상</li>
            </ul>
            
            <div class="tip-box">
                <i class="fas fa-lightbulb"></i> <strong>팁:</strong> 이미지의 크기가 작을수록 처리 속도가 빠릅니다. 얼굴이 포함된 이미지에서 가장 좋은 결과를 얻을 수 있습니다.
            </div>
        </p>
    </div>
</div>
"""

# UI 구성
with gr.Blocks(css=custom_css) as demo:
    gr.HTML(fontawesome_link)
    gr.HTML(header_html)
    
    with gr.Tabs(elem_classes="custom-tabs") as tabs:
        # 사용 설명서 탭을 첫 번째로 배치
        with gr.TabItem("📚 사용 설명서", elem_classes="tab-content"):
            gr.HTML("""
            <div class="user-manual">
                <div class="manual-title"><i class="fas fa-book"></i> 사용 설명서</div>
                
                <div class="manual-section">
                    <div class="manual-section-title"><i class="fas fa-magic"></i> 이커머스전용 이미지 생성기</div>
                    <p class="manual-text">
                        <span style="font-size: 1.1rem; font-weight: 600; color: #5561e9;">1️⃣ 이미지 업로드</span><br>
                        최대 3개의 이미지를 업로드할 수 있습니다. 첫 번째 이미지는 필수입니다.<br><br>
                        
                        <span style="font-size: 1.1rem; font-weight: 600; color: #5561e9;">2️⃣ 프롬프트 입력 또는 버튼 선택</span><br>
                        변환하고 싶은 내용을 직접 입력하거나, 미리 정의된 버튼 중 하나를 선택하세요.<br><br>
                        
                        <span style="font-size: 1.1rem; font-weight: 600; color: #5561e9;">3️⃣ 이미지 생성</span><br>
                        '이미지 생성 (1장)' 버튼을 클릭하면 1개의 이미지가, '이미지 생성 (4장)' 버튼을 클릭하면 4가지 버전의 이미지가 생성됩니다.
                    </p>
                </div>
                
                <div class="manual-section">
                    <div class="manual-section-title"><i class="fas fa-hashtag"></i> 이미지 참조 방법</div>
                    <p class="manual-text">
                        프롬프트에서 다음 태그를 사용하여 각 이미지를 참조할 수 있습니다:
                        <ul>
                            <li><strong style="font-size: 1.05rem;">#1</strong> - 첫 번째 이미지</li>
                            <li><strong style="font-size: 1.05rem;">#2</strong> - 두 번째 이미지</li>
                            <li><strong style="font-size: 1.05rem;">#3</strong> - 세 번째 이미지</li>
                        </ul>
                        
                        <div class="tip-box">
                            <i class="fas fa-lightbulb"></i> <strong>팁:</strong> 예를 들어, "(#1의 여성모델)이 (#2의 선글라스)를 착용한 모습"과 같이 사용할 수 있습니다.
                        </div>
                    </p>
                </div>
                
                <div class="manual-section">
                    <div class="manual-section-title"><i class="fas fa-tools"></i> 주요 기능</div>
                    <div class="manual-text">
                        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)); gap: 1rem;">
                            <div>
                                <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-exchange-alt"></i> 이미지 변경</p>
                                포즈나 스타일을 변경하면서 원본의 주요 요소는 유지합니다.
                            </div>
                            
                            <div>
                                <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-eraser"></i> 글자 지우기/변경</p>
                                이미지의 텍스트를 자연스럽게 제거하거나 다른 텍스트로 변경합니다.
                            </div>
                            
                            <div>
                                <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-tshirt"></i> 가상 상품착용</p>
                                모델에게 다른 이미지의 상품을 자연스럽게 착용시킵니다.
                            </div>
                            
                            <div>
                                <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-image"></i> 배경 바꾸기</p>
                                주요 피사체는 유지하며 배경만 교체합니다.
                            </div>
                            
                            <div>
                                <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-cut"></i> 부분 지우기</p>
                                이미지의 특정 부분을 지우고 배경과 자연스럽게 조화되도록 채웁니다.
                            </div>
                            
                            <div>
                                <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-wine-glass"></i> 상품들고 있기</p>
                                모델이 상품을 들고 있는 자연스러운 포즈의 이미지를 생성합니다.
                            </div>
                            
                            <div>
                                <p style="font-size: 1.1rem; font-weight: 600; color: #ff6b6b;"><i class="fas fa-expand-alt"></i> 이미지 확장</p>
                                원본 이미지 주변에 자연스럽게 추가 영역을 생성하여 이미지를 확장합니다.
                            </div>
                        </div>
                    </div>
                </div>
                
                <div class="manual-section">
                    <div class="manual-section-title"><i class="fas fa-search-plus"></i> 이미지 업스케일러</div>
                    <p class="manual-text">
                        <span style="font-size: 1.1rem; font-weight: 600; color: #5561e9;">이미지 개선 기능</span><br>
                        저해상도 이미지를 고해상도로 변환하고, 얼굴을 자연스럽게 복원합니다.<br>
                        특히 다음과 같은 경우에 효과적입니다:
                        <ul>
                            <li>해상도가 낮은 이미지 개선</li>
                            <li>흐릿하거나 노이즈가 있는 얼굴 복원</li>
                            <li>이미지 전반적인 선명도 향상</li>
                        </ul>
                        
                        <div class="tip-box">
                            <i class="fas fa-lightbulb"></i> <strong>팁:</strong> 이미지의 크기가 작을수록 처리 속도가 빠릅니다. 얼굴이 포함된 이미지에서 가장 좋은 결과를 얻을 수 있습니다.
                        </div>
                    </p>
                </div>
                
                <div class="manual-section">
                    <div class="manual-section-title"><i class="fas fa-cut"></i> 배경 지우기 기능</div>
                    <p class="manual-text">
                        <span style="font-size: 1.1rem; font-weight: 600; color: #5561e9;">자동 배경 제거</span><br>
                        이미지에서 배경을 자동으로 감지하고 제거하여 투명한 PNG 이미지로 변환합니다.<br>
                        주요 사용 사례:
                        <ul>
                            <li>제품 이미지의 배경 제거</li>
                            <li>인물 사진에서 배경 분리</li>
                            <li>여러 이미지 합성을 위한 준비</li>
                        </ul>
                        
                        <div class="tip-box">
                            <i class="fas fa-lightbulb"></i> <strong>팁:</strong> 배경과 주요 피사체 간의 대비가 뚜렷할수록 더 정확한 결과를 얻을 수 있습니다. 복잡한 배경의 경우 미세 조정이 필요할 수 있습니다.
                        </div>
                    </p>
                </div>
                
                <div class="manual-section">
                    <div class="manual-section-title"><i class="fas fa-question-circle"></i> 자주 묻는 질문</div>
                    <p class="manual-text">
                        <strong style="font-size: 1.1rem; color: #5561e9;">Q: 어떤 형식의 이미지를 업로드할 수 있나요?</strong><br>
                        A: JPG, PNG, WEBP 등 대부분의 일반적인 이미지 형식을 지원합니다.<br><br>
                        
                        <strong style="font-size: 1.1rem; color: #5561e9;">Q: 이미지 생성 및 처리에 얼마나 시간이 걸리나요?</strong><br>
                        A: 일반적으로 10-30초 정도 소요되며, 이미지 크기와 복잡도에 따라 달라질 수 있습니다.<br><br>
                        
                        <strong style="font-size: 1.1rem; color: #5561e9;">Q: 생성된 이미지를 어떻게 저장할 수 있나요?</strong><br>
                        A: 생성된 이미지에 마우스 오른쪽 버튼을 클릭하여 '이미지 저장' 옵션을 선택하면 됩니다.
                    </p>
                </div>
            </div>
            """)
            
        with gr.TabItem("✨ 이커머스 이미지 생성기", elem_classes="tab-content"):
            with gr.Row(equal_height=True):
                with gr.Column(scale=1):
                    with gr.Group(elem_classes="panel"):
                        gr.HTML('<div class="section-title"><i class="fas fa-upload"></i> 이미지 업로드</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)
                        
                        gr.HTML('<div class="section-title"><i class="fas fa-keyboard"></i> 프롬프트 입력</div>')
                        prompt_input = gr.Textbox(
                            lines=3,
                            placeholder="프롬프트를 입력하거나 비워두면 자동 합성됩니다. '#1', '#2', '#3'으로 각 이미지를 참조할 수 있습니다.",
                            label="프롬프트 (선택 사항)",
                            elem_classes="custom-input"
                        )
                        
                        gr.HTML('<div class="section-title"><i class="fas fa-sliders-h"></i> 변환 옵션</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")
                            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")
                        
                        submit_single_btn = gr.Button('✨ 이미지 생성 (1장)', elem_classes="custom-button primary")
                        submit_btn = gr.Button('✨ 이미지 생성 (4장)', elem_classes="custom-button primary")
                
                with gr.Column(scale=1):
                    with gr.Group(elem_classes="panel"):
                        gr.HTML('<div class="section-title"><i class="fas fa-images"></i> 생성된 이미지</div>')
                        with gr.Row():
                            with gr.Column():
                                output_image1 = gr.Image(label="이미지 #1", elem_classes="image-container", height=400)
                                output_image3 = gr.Image(label="이미지 #3", elem_classes="image-container", height=400)
                            with gr.Column():
                                output_image2 = gr.Image(label="이미지 #2", elem_classes="image-container", height=400)
                                output_image4 = gr.Image(label="이미지 #4", elem_classes="image-container", height=400)
                        
                        gr.HTML('<div class="section-title"><i class="fas fa-info-circle"></i> 결과 정보</div>')
                        output_text = gr.Textbox(label="상태 메시지", lines=2, elem_classes="custom-input")
                        prompt_display = gr.Textbox(label="사용된 프롬프트 (영어)", visible=True, lines=2, elem_classes="custom-input")
            
            gr.HTML('<div class="section-title"><i class="fas fa-lightbulb"></i> 예제 이미지</div>')
            with gr.Group(elem_classes="panel"):
                examples = [
                    ["down/모델.jpg", None, None, "(#1의 여성)이 살짝 뒤로 돌아보는 모습으로 최대한 이전 seed를 유지한테 자연스럽게 변경하라."],
                    ["down/상어레고모형.png", None, None, "(#1 레모모형)에서 청색상어레고만 검은색 고래레고로 변경하고 나머지 부분은 seed를 변경을 하지마라."],
                    ["down/중국어.png", None, None, "(#1 이미지)에 있는 중국어를 모두 제거하라."],
                    ["down/텍스트.webp", None, None, '(#1의 텍스트)를 스타일을 유지한체 텍스트만 "Hello"로 바꿔라'],
                    ["down/모델.jpg", "down/선글라스.png", "down/청바지.png", "(#1의 여성모델)이 신체 비율과 포즈는 유치한 체 (#2의 선글라스)와 (#3의 청바지)를 직접 모델이 착용한것 처럼 자연스럽게 이미지를 생성하라."],
                    ["down/모델.jpg", "down/선글라스.png", "down/카페전경.png", "(#1의 여성모델)이 신체 비율과 포즈는 유치한 체 (#2의 선글라스)를 직접 모델이 착용한 것처럼 (#3의 장소)에서 의자에 앉아 있는 자연스러운 이미지를 생성하라."],
                    ["down/모델.jpg", "down/와인잔.png", None, "(#1여성모델)이 (#2 선글라스)를 착용하고 (#3뒷 배경의 카페전체가 보이며) 그 앞 의자에 앉아 있는 모습"],
                    ["down/모델.jpg", "down/카페전경.png", None, "(#1의 여성모델)이 (#2 이미지의 배경)을 주요 피사체는 그대로 유지하여 이미지의 분위기가 자연스럽게 어우러지도록 생성하라."],
                    ["down/상어레고모형.png", None, None, "(#1의 레고모형)에서 청색상어레고를 제거한 후, 그 자리를 주변 배경과 자연스럽게 어우러지도록 채워주세요. 단, 이미지의 다른 부분의 주요 요소는 동일하게 유지해 해야한다."],
                    ["down/카페전경.png", None, None, "(#1 이미지)를 원본그대로 중앙에 두고 비율로 유지한체 확장하라.(outpaint)"]
                ]
                
                gr.Examples(
                    examples=examples,
                    inputs=[image1_input, image2_input, image3_input, prompt_input]
                )
            
        with gr.TabItem("🔍 이미지 업스케일러", elem_classes="tab-content"):
            with gr.Row():
                with gr.Column(elem_classes="panel"):
                    gr.HTML('<div class="section-title"><i class="fas fa-search-plus"></i> 이미지 업스케일 & 복원</div>')
                    gr.HTML("""
                    <p style="margin-bottom: 1rem; font-size: 1.05rem;">
                         이미지를 업스케일러합니다.
                    </p>
                    <div class="tip-box">
                        <i class="fas fa-lightbulb"></i> <strong>팁:</strong> 해상도가 낮거나 품질이 떨어지는 이미지, 특히 얼굴 이미지에 효과적입니다.
                    </div>
                    """)
                    
                    with gr.Row():
                        with gr.Column():
                            gr.HTML('<div class="section-title"><i class="fas fa-upload"></i> 입력 이미지</div>')
                            gfpgan_input = gr.Image(type="filepath", label="업로드", elem_classes="image-container")
                        
                        with gr.Column():
                            gr.HTML('<div class="section-title"><i class="fas fa-download"></i> 결과 이미지</div>')
                            gfpgan_output = gr.Image(type="numpy", label="업스케일 결과", elem_classes="image-container")
                    
                    gfpgan_btn = gr.Button('✨ 업스케일 및 복원', elem_classes="custom-button primary")
        
        # 배경 지우기 탭
        with gr.TabItem("✂️ 배경 지우기", elem_classes="tab-content"):
            with gr.Row():
                with gr.Column(elem_classes="panel"):
                    gr.HTML('<div class="section-title"><i class="fas fa-cut"></i> 이미지 배경 제거</div>')
                    gr.HTML("""
                    <p style="margin-bottom: 1rem; font-size: 1.05rem;">
                         이미지에서 배경을 자동으로 제거하여 투명한 PNG 이미지로 변환합니다.
                    </p>
                    <div class="tip-box">
                        <i class="fas fa-lightbulb"></i> <strong>팁:</strong> 배경과 주요 피사체 간의 대비가 뚜렷할수록 더 나은 결과를 얻을 수 있습니다.
                    </div>
                    """)
                    
                    with gr.Row():
                        with gr.Column():
                            gr.HTML('<div class="section-title"><i class="fas fa-upload"></i> 입력 이미지</div>')
                            bg_remove_input = gr.Image(type="pil", label="업로드", elem_classes="image-container")
                        
                        with gr.Column():
                            gr.HTML('<div class="section-title"><i class="fas fa-download"></i> 결과 이미지</div>')
                            bg_remove_output = gr.Image(type="pil", label="배경 제거 결과", elem_classes="image-container", format="png")
                    
                    with gr.Row():
                        bg_remove_status = gr.Textbox(label="상태", lines=2, elem_classes="custom-input")
                        bg_remove_btn = gr.Button('✂️ 배경 제거하기', elem_classes="custom-button primary")
    
    # 버튼 이벤트 연결
    image_change_btn1.click(
        fn=lambda: "(#1의 여성)이 살짝 뒤로 돌아보는 모습으로 최대한 이전 seed를 유지한테 자연스럽게 변경하라.",
        inputs=[],
        outputs=prompt_input
    )
    image_change_btn2.click(
        fn=lambda: "(#1 레모모형)에서 청색상어레고만 검은색 고래레고로 변경하고 나머지 부분은 seed를 변경을 하지마라.",
        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 이미지)를 원본그대로 중앙에 두고 비율로 유지한체 확장하라.(outpaint)",
        inputs=[],
        outputs=prompt_input
    )
    
    # 단일 이미지 생성 버튼 이벤트 연결
    submit_single_btn.click(
        fn=lambda image1, image2, image3, prompt: process_images_with_prompt(image1, image2, image3, prompt, 0),
        inputs=[image1_input, image2_input, image3_input, prompt_input],
        outputs=[output_image1, output_text, prompt_display],
    )
    
    # 4장 이미지 생성 버튼 이벤트 연결
    submit_btn.click(
        fn=generate_multiple_images,
        inputs=[image1_input, image2_input, image3_input, prompt_input],
        outputs=[output_image1, output_image2, output_image3, output_image4, output_text, prompt_display],
    )
    
    gfpgan_btn.click(
        fn=upscaler_korean, 
        inputs=gfpgan_input, 
        outputs=gfpgan_output
    )
    
    # 배경 제거 버튼 이벤트 연결
    bg_remove_btn.click(
        fn=remove_background,
        inputs=bg_remove_input,
        outputs=[bg_remove_output, bg_remove_status]
    )

demo.queue()
demo.launch()