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

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_KEYS = []  # API 키 목록
current_key_index = 0  # 현재 사용 중인 키 인덱스

def initialize_api_keys():
    """API 키 목록을 초기화하는 함수"""
    global API_KEYS
    # 환경 변수에서 API 키 가져오기
    key1 = os.environ.get("GEMINI_API_KEY_1", "")
    key2 = os.environ.get("GEMINI_API_KEY_2", "")
    key3 = os.environ.get("GEMINI_API_KEY_3", "")
    key4 = os.environ.get("GEMINI_API_KEY_4", "")
    key5 = os.environ.get("GEMINI_API_KEY_5", "")
    
    # 빈 문자열이 아닌 키만 추가
    if key1:
        API_KEYS.append(key1)
    if key2:
        API_KEYS.append(key2)
    if key3:
        API_KEYS.append(key3)
    if key4:
        API_KEYS.append(key4)
    if key5:
        API_KEYS.append(key5)
    
    # 기존 GEMINI_API_KEY가 있으면 추가
    default_key = os.environ.get("GEMINI_API_KEY", "")
    if default_key and default_key not in API_KEYS:
        API_KEYS.append(default_key)
    
    logger.info(f"API 키 {len(API_KEYS)}개가 로드되었습니다.")

def get_next_api_key():
    """다음 API 키를 가져오는 함수"""
    global current_key_index
    
    if not API_KEYS:
        return None
    
    # 현재 키 가져오기
    api_key = API_KEYS[current_key_index]
    
    # 다음 키 인덱스로 업데이트
    current_key_index = (current_key_index + 1) % len(API_KEYS)
    
    return api_key

# ========== 이미지 생성기 관련 함수 ==========
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 키가 설정되지 않았습니다. 환경 변수에 GEMINI_API_KEY_1, GEMINI_API_KEY_2, GEMINI_API_KEY_3, GEMINI_API_KEY_4, GEMINI_API_KEY_5 중 하나 이상을 설정해주세요."

        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

# ========== 이미지 편집기 관련 함수 ==========
def adjust_brightness(image, value):
    """이미지 밝기 조절"""
    value = float(value - 1) * 100  # 0-2 범위를 -100에서 +100으로 변환
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    v = cv2.add(v, value)
    v = np.clip(v, 0, 255)
    final_hsv = cv2.merge((h, s, v))
    return cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)

def adjust_contrast(image, value):
    """이미지 대비 조절"""
    value = float(value)
    return np.clip(image * value, 0, 255).astype(np.uint8)

def adjust_saturation(image, value):
    """이미지 채도 조절"""
    value = float(value - 1) * 100  # 0-2 범위를 -100에서 +100으로 변환
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    s = cv2.add(s, value)
    s = np.clip(s, 0, 255)
    final_hsv = cv2.merge((h, s, v))
    return cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)

def adjust_temperature(image, value):
    """이미지 색온도 조절 (색상 밸런스)"""
    value = float(value) * 30  # 효과 스케일 조절
    b, g, r = cv2.split(image)
    if value > 0:  # 따뜻하게
        r = cv2.add(r, value)
        b = cv2.subtract(b, value)
    else:  # 차갑게
        r = cv2.add(r, value)
        b = cv2.subtract(b, value)
    
    r = np.clip(r, 0, 255)
    b = np.clip(b, 0, 255)
    return cv2.merge([b, g, r])

def adjust_tint(image, value):
    """이미지 색조 조절"""
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv_image)
    h = cv2.add(h, int(value))
    h = np.clip(h, 0, 179)  # Hue 값은 0-179 범위
    final_hsv = cv2.merge((h, s, v))
    return cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)

def adjust_exposure(image, value):
    """이미지 노출 조절"""
    enhancer = ImageEnhance.Brightness(Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)))
    img_enhanced = enhancer.enhance(1 + float(value) / 5.0)
    return cv2.cvtColor(np.array(img_enhanced), cv2.COLOR_RGB2BGR)

def adjust_vibrance(image, value):
    """이미지 활기 조절"""
    img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    converter = ImageEnhance.Color(img)
    factor = 1 + (float(value) / 100.0)
    img = converter.enhance(factor)
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

def adjust_color_mixer_blues(image, value):
    """이미지 컬러 믹서 (블루) 조절"""
    b, g, r = cv2.split(image)
    b = cv2.add(b, float(value))
    b = np.clip(b, 0, 255)
    return cv2.merge([b, g, r])

def adjust_shadows(image, value):
    """이미지 그림자 조절"""
    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    enhancer = ImageEnhance.Brightness(pil_image)
    factor = 1 + (float(value) / 100.0)
    pil_image = enhancer.enhance(factor)
    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_BGR2RGB)

def process_image(image, brightness, contrast, saturation, temperature, tint, exposure, vibrance, color_mixer_blues, shadows):
    """모든 조정 사항을 이미지에 적용"""
    if image is None:
        return None
        
    # PIL 이미지를 OpenCV 형식으로 변환
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    
    # 조정 사항 순차 적용
    image = adjust_brightness(image, brightness)
    image = adjust_contrast(image, contrast)
    image = adjust_saturation(image, saturation)
    image = adjust_temperature(image, temperature)
    image = adjust_tint(image, tint)
    image = adjust_exposure(image, exposure)
    image = adjust_vibrance(image, vibrance)
    image = adjust_color_mixer_blues(image, color_mixer_blues)
    image = adjust_shadows(image, shadows)
    
    # PIL 이미지로 다시 변환
    return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

def download_edited_image(image, input_image_name):
    """이미지를 JPG 형식으로 저장하고 경로 반환"""
    if image is None:
        return None

    # 한국 시간 타임스탬프 생성
    def get_korean_timestamp():
        korea_time = datetime.utcnow() + timedelta(hours=9)
        return korea_time.strftime('%Y%m%d_%H%M%S')

    timestamp = get_korean_timestamp()
    if input_image_name and hasattr(input_image_name, 'name'):
        base_name = input_image_name.name.split('.')[0]  # 파일 객체에서 이름 추출
    else:
        base_name = "이미지"

    file_name = f"[끝장AI]끝장필터_{base_name}_{timestamp}.jpg"

    # 파일 저장
    temp_file_path = tempfile.gettempdir() + "/" + file_name
    image.save(temp_file_path, format="JPEG")
    return temp_file_path

# 커스텀 CSS 스타일 - 새로운 스타일 적용
custom_css = """
:root {
    --primary-color: #FB7F0D;
    --secondary-color: #ff9a8b;
    --accent-color: #FF6B6B;
    --background-color: #FFF3E9;
    --card-bg: #ffffff;
    --text-color: #334155;
    --border-radius: 18px;
    --shadow: 0 8px 30px rgba(251, 127, 13, 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;  /* 200%에서 100%로 변경 */
    width: 100% !important;      /* 추가: 너비 100% 지정 */
    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;
}

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

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

/* 섹션 제목 */
.section-title {
    font-size: 22px !important;
    font-weight: 700 !important;
    color: #333333 !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);
}

/* 버튼 스타일링 */
.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;
}

/* 이미지 컨테이너 */
.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);
}

/* 탭 스타일 개선 */
.tabs {
    border-bottom: none !important;
}

.tab-nav {
    background-color: transparent !important;
    border-bottom: 1px solid #eeeeee !important;
    padding: 0 !important;
}

.tab-nav button {
    border-radius: var(--border-radius) var(--border-radius) 0 0 !important;
    margin-right: 5px !important;
    padding: 12px 20px !important;
    font-size: 18px !important;
    font-weight: 600 !important;
    border: 1px solid #eeeeee !important;
    border-bottom: none !important;
    background-color: rgba(255, 255, 255, 0.7) !important;
    color: var(--text-color) !important;
    transition: all 0.3s ease !important;
    min-width: 150px !important;
    text-align: center !important;
}

.tab-nav button.selected {
    background-color: var(--primary-color) !important;
    color: white !important;
    border-color: var(--primary-color) !important;
    box-shadow: 0 -2px 6px rgba(251, 127, 13, 0.2) !important;
}

.tab-nav button:hover:not(.selected) {
    background-color: var(--background-color) !important;
    border-bottom: none !important;
}

/* 입력 필드 스타일 */
.gr-input, .gr-text-input, .gr-sample-inputs {
    border-radius: var(--border-radius) !important;
    border: 1px solid #dddddd !important;
    padding: 12px !important;
    box-shadow: inset 0 1px 3px rgba(0, 0, 0, 0.05) !important;
    transition: all 0.3s ease !important;
}

.gr-input:focus, .gr-text-input:focus {
    border-color: var(--primary-color) !important;
    outline: none !important;
    box-shadow: 0 0 0 2px rgba(251, 127, 13, 0.2) !important;
}

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

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

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

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

/* 편집기 특별 스타일 */
.editor-section {
    background-color: var(--card-bg);
    border-radius: var(--border-radius);
    box-shadow: var(--shadow);
    padding: 1.5rem;
    margin-bottom: 1.5rem;
}

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

.editor-title i {
    margin-right: 0.5rem;
}

.download-button {
    background-color: var(--accent-color) !important;
    color: white !important;
    border: none !important;
    padding: 10px !important;
    font-size: 12px !important;
    border-radius: var(--border-radius) !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
}

.download-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.15) !important;
}

.download-container {
    display: flex;
    flex-direction: column;
    align-items: center;
    width: 100%;
}

.download-output {
    width: 100%;
    margin-top: 1rem;
}

/* 애니메이션 스타일 */
@keyframes fadeIn {
    from { opacity: 0; transform: translateY(10px); }
    to { opacity: 1; transform: translateY(0); }
}

.fade-in {
    animation: fadeIn 0.5s ease-out;
}

/* Examples 섹션 스타일 */
.examples-section {
    display: grid;
    grid-template-columns: repeat(auto-fill, minmax(300px, 1fr));
    gap: 1.5rem;
    margin-top: 1rem;
}

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

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

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

    .examples-section {
        grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
    }
}
"""  # CSS 문자열 종료

# 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 = ""
image_editor_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.Tabs(elem_classes="tabs") as tabs:
        # 이미지 생성기 탭
        with gr.Tab("✨ 이미지 생성기", elem_classes="tab-content"):
            # 사용 가이드 섹션 제거
            # gr.HTML(image_generator_guide_html)
            with gr.Row(equal_height=True):
                with gr.Column(scale=1):
                    # API 키 입력 섹션 제거
                    
                    # ======== 이미지 업로드 및 설정 섹션 ========
                    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('✨ 이미지 생성 (4장)', 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)
                                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)
                    
                    # ======== 결과 정보 섹션 ========
                    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)  # 모든 예제를 한 페이지에 표시
            )
        
        # 이미지 편집기 탭 추가
        with gr.Tab("🎨 이미지 편집기", elem_classes="tab-content"):
            # 사용 가이드 섹션 제거
            # gr.HTML(image_editor_guide_html)
            
            with gr.Row():
                # 왼쪽 열: 비율 1
                with gr.Column(scale=1):
                    # ======== 이미지 업로드 섹션 ========
                    with gr.Group():
                        gr.HTML('<div class="section-title"><i class="fas fa-upload"></i> <span>이미지 업로드</span></div>')
                        edit_input_image = gr.Image(type="pil", label="편집할 이미지", elem_classes="image-container")

                    # ======== 이미지 조정 섹션 ========
                    with gr.Group():
                        gr.HTML('<div class="section-title"><i class="fas fa-sliders-h"></i> <span>이미지 조정</span></div>')
                        brightness_slider = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="밝기 조절")
                        contrast_slider = gr.Slider(0.5, 1.5, value=1.0, step=0.1, label="대비 조절")
                        saturation_slider = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="채도 조절")
                        temperature_slider = gr.Slider(-1.0, 1.0, value=0.0, step=0.1, label="색온도 조절")
                        tint_slider = gr.Slider(-100, 100, value=0, step=1, label="색조 조절")
                        exposure_slider = gr.Slider(-5.0, 5.0, value=0.0, step=0.1, label="노출 조절")
                        vibrance_slider = gr.Slider(-100.0, 100.0, value=0.0, step=1.0, label="활기 조절")
                        color_mixer_blues_slider = gr.Slider(-100.0, 100.0, value=0.0, step=1.0, label="컬러 믹서 (블루)")
                        shadows_slider = gr.Slider(-100.0, 100.0, value=0.0, step=1.0, label="그림자 조절")

                # 오른쪽 열: 비율 1
                with gr.Column(scale=1):
                    # ======== 편집된 이미지 섹션 ========
                    with gr.Group():
                        gr.HTML('<div class="section-title"><i class="fas fa-images"></i> <span>편집된 이미지</span></div>')
                        edit_output_image = gr.Image(type="pil", label="편집된 이미지", elem_classes="image-container")
                    
                    # ======== 저장 섹션 ========
                    with gr.Group():
                        gr.HTML('<div class="section-title"><i class="fas fa-download"></i> <span>저장</span></div>')
                        with gr.Row(elem_classes="download-container"):
                            download_button = gr.Button("JPG로 변환하기", elem_classes="download-button")
                        with gr.Row(elem_classes="download-container"):
                            download_output = gr.File(label="JPG 이미지 다운로드", elem_classes="download-output")
    
    # ========== 이미지 생성기 이벤트 연결 ==========
    # 버튼 이벤트 연결
    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],
    )
    
    # 4장 이미지 생성 버튼 이벤트 연결 - API 키 입력값 제거
    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],
    )
    
    # ========== 이미지 편집기 이벤트 연결 ==========
    # 이미지 처리 함수 연결
    edit_inputs = [
        edit_input_image,
        brightness_slider,
        contrast_slider,
        saturation_slider,
        temperature_slider,
        tint_slider,
        exposure_slider,
        vibrance_slider,
        color_mixer_blues_slider,
        shadows_slider
    ]

    edit_input_components = [
        brightness_slider,
        contrast_slider,
        saturation_slider,
        temperature_slider,
        tint_slider,
        exposure_slider,
        vibrance_slider,
        color_mixer_blues_slider,
        shadows_slider
    ]

    for input_component in edit_input_components:
        input_component.change(
            fn=process_image,
            inputs=edit_inputs,
            outputs=edit_output_image
        )
    
    # 이미지 업로드 시 자동으로 편집 업데이트
    edit_input_image.change(
        fn=process_image,
        inputs=edit_inputs,
        outputs=edit_output_image
    )

    # 다운로드 버튼 기능
    download_button.click(
        fn=download_edited_image,
        inputs=[edit_output_image, edit_input_image],
        outputs=download_output
    )

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