import os
import tempfile
from PIL import Image
import gradio as gr
import logging
import re
from io import BytesIO

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__)

def save_binary_file(file_name, data):
    with open(file_name, "wb") as f:
        f.write(data)

def translate_prompt_to_english(prompt):
    """
    입력된 프롬프트에 한글이 포함되어 있으면 Gemini‑2.0‑flash 모델을 사용하여 영어로 번역합니다.
    한글이 없으면 원본 프롬프트를 그대로 반환합니다.
    """
    if not re.search("[가-힣]", prompt):
        return prompt
    try:
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            logger.error("Gemini API 키가 설정되지 않았습니다.")
            return prompt
        client = genai.Client(api_key=api_key)
        translation_prompt = f"Translate the following Korean text to English:\n\n{prompt}"
        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():
            logger.info(f"Translated text: {translated_text.strip()}")
            return translated_text.strip()
        else:
            logger.warning("번역 결과가 없습니다. 원본 프롬프트 사용")
            return prompt
    except Exception as e:
        logger.exception("번역 중 오류 발생:")
        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):
    """
    API 호출을 통해 이미지를 생성하고 결과 이미지를 반환합니다.
    """
    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}")

        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):
    """
    3개의 이미지와 프롬프트를 처리하여 최종 영어 프롬프트(final_prompt)를 생성한 후,
    API를 호출하여 결과 이미지를 반환합니다.
    """
    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."
                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."
                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."
                logger.info("Default prompt generated for three images")

        result_img, status = generate_with_images(final_prompt, valid_images)
        return result_img, status, final_prompt
    except Exception as e:
        logger.exception("이미지 처리 중 오류 발생:")
        return None, f"오류 발생: {str(e)}", prompt

def process_and_show_prompt(image1, image2, image3, prompt):
    try:
        result_img, status, final_prompt = process_images_with_prompt(image1, image2, image3, prompt)
        return result_img, status, final_prompt
    except Exception as e:
        logger.exception("처리 중 오류 발생:")
        return None, f"오류 발생: {str(e)}", prompt

# 예제 실행 함수들
def run_example_1():
    """예제 1: 이미지 변경 예제"""
    input_path = os.path.join("down", "1_in-1.png")
    output_path = os.path.join("down", "1_out-1.webp")
    try:
        input_img = Image.open(input_path)
    except Exception as e:
        return None, f"입력 이미지 열기 오류: {str(e)}", ""
    prompt = "#1 이미지의 [청색 상어레고를 검은 고래레고]으로 변경하라."
    result_img, status, final_prompt = process_and_show_prompt(input_img, None, None, prompt)
    if result_img is not None:
        result_img.save(output_path, "WEBP")
    return result_img, status, final_prompt

def run_example_2():
    """예제 2: 글자지우기 예제"""
    input_path = os.path.join("down", "2_in-1.png")
    output_path = os.path.join("down", "2_out-1.webp")
    try:
        input_img = Image.open(input_path)
    except Exception as e:
        return None, f"입력 이미지 열기 오류: {str(e)}", ""
    prompt = "#1 이미지의 [중국어를 모두]를 제거하라."
    result_img, status, final_prompt = process_and_show_prompt(input_img, None, None, prompt)
    if result_img is not None:
        result_img.save(output_path, "WEBP")
    return result_img, status, final_prompt


# 기존 Gradio 인터페이스 구성 (예제 적용 외 다른 코드는 그대로 유지)
with gr.Blocks() as demo:
    gr.HTML(
        """
        <div style="text-align: center; margin-bottom: 1rem;">
            <h1>Gemini for Image Editing</h1>
            <p>Upload an image and enter a prompt to generate outputs.</p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                type="pil",
                label="Upload Image",
                image_mode="RGBA",
                elem_id="image-input"
            )
            gemini_api_key = gr.Textbox(
                lines=1,
                placeholder="Enter Gemini API Key (optional)",
                label="Gemini API Key (optional)"
            )
            prompt_input = gr.Textbox(
                lines=2,
                placeholder="Enter prompt here...",
                label="Prompt"
            )
            submit_btn = gr.Button("Generate")
        with gr.Column():
            output_gallery = gr.Gallery(label="Generated Outputs")
            output_text = gr.Textbox(
                label="Gemini Output",
                placeholder="Text response will appear here if no image is generated."
            )
    
    gr.Markdown("## Try these examples", elem_classes="gr-examples-header")
    
    examples = [
        ["data/1.webp", 'change text to "AMEER"', ""],
        ["data/2.webp", "remove the spoon from hand only", ""],
        ["data/3.webp", 'change text to "Make it "', ""],
        ["data/1.jpg", "add joker style only on face", ""],
        ["data/1777043.jpg", "add joker style only on face", ""],
        ["data/2807615.jpg", "add lipstick on lip only", ""],
        ["data/76860.jpg", "add lipstick on lip only", ""],
        ["data/2807615.jpg", "make it happy looking face only", ""],
    ]
    
    gr.Examples(
        examples=examples,
        inputs=[image_input, prompt_input],
        elem_id="examples-grid"
    )

demo.queue(max_size=50).launch()