import os
import tempfile
from PIL import Image
import gradio as gr
import logging
import re
import io
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 preprocess_prompt(prompt, image1, image2, image3):
    """
    프롬프트를 처리하고 기능 명령을 해석
    """
    # 이미지 없는 참조 확인 및 처리
    has_img1 = image1 is not None
    has_img2 = image2 is not None
    has_img3 = image3 is not None
    
    # #1, #2, #3 참조를 설명으로 변환 (이미지가 없는 경우 무시)
    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:
        # 기본 프롬프트 제공
        prompt = "첫 번째 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."
    
    elif "2. 글자지우기" in prompt:
        # 기본 프롬프트 제공
        prompt = "첫 번째 이미지에서 모든 텍스트를 찾아 자연스럽게 제거해주세요. 깔끔한 이미지로 만들어주세요."
    
    elif "3. 얼굴바꾸기" in prompt:
        prompt = "첫 번째 이미지의 인물 얼굴을 두 번째 이미지의 얼굴로 자연스럽게 교체해주세요. 얼굴의 표정과 특징은 두 번째 이미지를 따르되, 나머지 부분은 첫 번째 이미지를 유지해주세요."
    
    elif "4. 옷바꾸기" in prompt:
        # 여러 이미지 참조 처리
        if "#3" in prompt or "또는 #3" in prompt:
            prompt = "첫 번째 이미지의 인물 의상을 두 번째 또는 세 번째 이미지의 의상으로 자연스럽게 교체해주세요. 의상의 스타일과 색상은 참조 이미지를 따르되, 신체 비율과 포즈는 첫 번째 이미지를 유지해주세요."
        else:
            prompt = "첫 번째 이미지의 인물 의상을 두 번째 이미지의 의상으로 자연스럽게 교체해주세요. 의상의 스타일과 색상은 두 번째 이미지를 따르되, 신체 비율과 포즈는 첫 번째 이미지를 유지해주세요."
    
    elif "5. 배경바꾸기" in prompt:
        prompt = "첫 번째 이미지의 배경을 두 번째 이미지의 배경으로 자연스럽게 교체해주세요. 첫 번째 이미지의 주요 피사체는 유지하고, 두 번째 이미지의 배경과 조화롭게 합성해주세요."
    
    elif "6. 이미지 합성(상품포함)" in prompt:
        # 여러 이미지 참조 처리
        if "#3" in prompt or "또는 #3" in prompt:
            prompt = "첫 번째 이미지와 두 번째, 세 번째 이미지를 자연스럽게 합성해주세요. 모든 이미지의 주요 요소를 포함하고, 특히 상품이 잘 보이도록 조화롭게 통합해주세요."
        else:
            prompt = "첫 번째 이미지와 두 번째 이미지를 자연스럽게 합성해주세요. 두 이미지의 주요 요소를 포함하고, 특히 상품이 잘 보이도록 조화롭게 통합해주세요."
    
    elif "7. 이미지 합성(스타일적용)" in prompt:
        prompt = "첫 번째 이미지의 내용을 두 번째 이미지의 스타일로 변환해주세요. 첫 번째 이미지의 주요 피사체와 구도는 유지하되, 두 번째 이미지의 예술적 스타일, 색상, 질감을 적용해주세요."
    
    # 간단한 색상 변경 요청 처리
    elif "을 붉은색으로 바꿔라" in prompt or "를 붉은색으로 바꿔라" in prompt:
        prompt = "첫 번째 이미지를 붉은색 톤으로 변경해주세요. 전체적인 색상을 붉은 계열로 조정하고 자연스러운 느낌을 유지해주세요."
    
    # 명확한 이미지 생성 요청 추가
    prompt += " 이미지를 생성해주세요."
    
    return prompt

def generate_with_images(prompt, images, max_retries=2):
    """
    공식 문서에 기반한 올바른 API 호출 방식 구현
    실패 시 재시도 기능 추가
    """
    retries = 0
    last_error = None
    
    while retries <= max_retries:
        try:
            # API 키 확인
            api_key = os.environ.get("GEMINI_API_KEY")
            if not api_key:
                return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."
                
            # Gemini 클라이언트 초기화
            client = genai.Client(api_key=api_key)
            
            logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt} (시도: {retries+1}/{max_retries+1})")
            
            # 이미지 및 텍스트 컨텐츠 준비
            parts = []
            
            # 이미지 파트 추가
            for idx, img in enumerate(images, 1):
                if img is not None:
                    # 이미지를 바이트로 변환
                    buffered = BytesIO()
                    img.save(buffered, format="PNG")
                    image_bytes = buffered.getvalue()
                    
                    # 이미지 파트 생성
                    parts.append({
                        'inlineData': {
                            'mimeType': 'image/png',
                            'data': base64.b64encode(image_bytes).decode('utf-8')
                        }
                    })
                    logger.info(f"이미지 #{idx} 추가됨")
            
            # 텍스트 프롬프트 추가
            parts.append({
                'text': prompt
            })
            
            # 모델 생성 요청
            request = {
                'contents': [{
                    'parts': parts
                }],
                'generation_config': {
                    'temperature': 1.0,
                    'topP': 0.95,
                    'topK': 40,
                    'maxOutputTokens': 8192,
                    'responseMimeType': 'image/png'
                }
            }
            
            # 이미지 생성 API 호출
            model = client.get_model('gemini-pro-vision')
            response = model.generate_content(request)
            
            # 응답 유효성 확인
            if not response or not response.parts:
                if retries < max_retries:
                    retries += 1
                    logger.warning(f"유효한 응답을 받지 못했습니다. 재시도 중... ({retries}/{max_retries})")
                    continue
                return None, "이미지 생성에 실패했습니다. 유효한 응답을 받지 못했습니다."
            
            # 임시 파일 생성
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
                temp_path = tmp.name
                
                result_text = ""
                image_found = False
                
                # 이미지 데이터 확인 및 저장
                for part in response.parts:
                    if hasattr(part, 'mime_type') and part.mime_type.startswith('image/'):
                        # base64 디코딩 및 저장
                        img_data = base64.b64decode(part.data)
                        with open(temp_path, 'wb') as f:
                            f.write(img_data)
                        image_found = True
                        logger.info("응답에서 이미지 추출 성공")
                        break
                
                # 텍스트 추출 (옵션)
                if hasattr(response, 'text'):
                    result_text = response.text
                
                if not image_found:
                    if retries < max_retries:
                        retries += 1
                        logger.warning(f"API에서 이미지를 생성하지 못했습니다. 재시도 중... ({retries}/{max_retries})")
                        continue
                    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:
            last_error = str(e)
            logger.exception(f"이미지 생성 중 오류 발생 (시도: {retries+1}/{max_retries+1}):")
            if retries < max_retries:
                retries += 1
                logger.info(f"재시도 중... ({retries}/{max_retries})")
                continue
            return None, f"오류 발생: {last_error}"
    
    return None, f"최대 재시도 횟수 초과. 마지막 오류: {last_error}"


def process_images_with_prompt(image1, image2, image3, prompt):
    """
    3개의 이미지와 프롬프트를 처리하는 함수
    이미지 없이 프롬프트만으로도 동작
    """
    try:
        # 이미지 개수 확인
        images = [image1, image2, image3]
        valid_images = [img for img in images if img is not None]
        
        # 프롬프트 처리
        if not prompt or not prompt.strip():
            # 프롬프트가 없는 경우
            if not valid_images:
                # 이미지도 없고 프롬프트도 없는 경우 기본 프롬프트 사용
                prompt = "아름다운 풍경 이미지를 생성해주세요. 푸른 하늘, 초록색 언덕, 그리고 맑은 호수가 있는 평화로운 자연 풍경을 만들어주세요."
                logger.info("이미지와 프롬프트 모두 없음. 기본 프롬프트 생성")
            elif len(valid_images) == 1:
                prompt = "이 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."
                logger.info("단일 이미지 프롬프트 자동 생성")
            elif len(valid_images) == 2:
                prompt = "이 두 이미지를 자연스럽게 합성해주세요. 두 이미지의 요소를 조화롭게 통합하여 하나의 이미지로 만들어주세요."
                logger.info("두 이미지 합성 프롬프트 자동 생성")
            else:
                prompt = "이 세 이미지를 창의적으로 합성해주세요. 모든 이미지의 주요 요소를 포함하되 자연스럽고 일관된 하나의 장면으로 만들어주세요."
                logger.info("세 이미지 합성 프롬프트 자동 생성")
        else:
            # 프롬프트 전처리 및 기능 명령 해석
            prompt = preprocess_prompt(prompt, image1, image2, image3)
        
        # 이미지 생성
        return generate_with_images(prompt, valid_images)
            
    except Exception as e:
        logger.exception("이미지 처리 중 오류 발생:")
        return None, f"오류 발생: {str(e)}"

# 기능 선택 콜백
def update_prompt_from_function(function_choice):
    function_templates = {
        "1. 이미지 변경": "#1을 창의적으로 바꿔라",
        "2. 글자지우기": "#1에서 텍스트를 지워라",
        "3. 얼굴바꾸기": "#1의 인물을 #2의 얼굴로 바꿔라",
        "4. 옷바꾸기": "#1의 인물에 #2 또는 #3의 옷으로 변경하라",
        "5. 배경바꾸기": "#1의 이미지에 #2의 배경으로 자연스럽게 바꿔라",
        "6. 이미지 합성(상품포함)": "#1와 #2 또는 #3의 를 합성하라",
        "7. 이미지 합성(스타일적용)": "#1와 #2를 스타일로 변환하라"
    }
    
    return function_templates.get(function_choice, "")

# Gradio 인터페이스
with gr.Blocks() as demo:
    gr.HTML(
        """
        <div style="text-align: center; margin-bottom: 1rem;">
            <h1>간단한 이미지 생성기</h1>
            <p>이미지를 업로드하거나 프롬프트를 입력하여 이미지를 생성하세요. 둘 다 없어도 기본 이미지가 생성됩니다.</p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column():
            # 3개의 이미지 입력
            with gr.Row():
                image1_input = gr.Image(type="pil", label="#1", image_mode="RGB")
                image2_input = gr.Image(type="pil", label="#2", image_mode="RGB")
                image3_input = gr.Image(type="pil", label="#3", image_mode="RGB")
            
            # 기능 선택 드롭다운
            function_dropdown = gr.Dropdown(
                choices=[
                    "1. 이미지 변경", 
                    "2. 글자지우기", 
                    "3. 얼굴바꾸기", 
                    "4. 옷바꾸기",
                    "5. 배경바꾸기", 
                    "6. 이미지 합성(상품포함)", 
                    "7. 이미지 합성(스타일적용)"
                ],
                label="기능 선택",
                value=None
            )
            
            # 프롬프트 입력 (선택 사항)
            prompt_input = gr.Textbox(
                lines=3,
                placeholder="프롬프트를 입력하거나 비워두면 자동 합성됩니다. 이미지 없이도 프롬프트만으로 생성 가능합니다.",
                label="프롬프트 (선택 사항)"
            )
            
            # 생성 버튼
            submit_btn = gr.Button("이미지 생성", variant="primary")
        
        with gr.Column():
            # 결과 출력
            output_image = gr.Image(label="생성된 이미지")
            output_text = gr.Textbox(label="상태 메시지")
            
            # 사용된 프롬프트 표시
            prompt_display = gr.Textbox(label="사용된 프롬프트", visible=True)
    
    # 기능 선택 드롭다운 이벤트 (즉시 프롬프트 업데이트)
    function_dropdown.change(
        fn=update_prompt_from_function,
        inputs=[function_dropdown],
        outputs=[prompt_input]
    )
    
    # 이미지 생성 버튼 클릭 이벤트
    def process_and_show_prompt(image1, image2, image3, prompt):
        # 이미지 개수 확인
        images = [image1, image2, image3]
        valid_images = [img for img in images if img is not None]
        
        try:
            # 자동 프롬프트 생성 또는 프롬프트 전처리
            auto_prompt = prompt
            if not prompt or not prompt.strip():
                if not valid_images:
                    # 이미지도 없고 프롬프트도 없는 경우
                    auto_prompt = "아름다운 풍경 이미지를 생성해주세요. 푸른 하늘, 초록색 언덕, 그리고 맑은 호수가 있는 평화로운 자연 풍경을 만들어주세요."
                elif len(valid_images) == 1:
                    auto_prompt = "이 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."
                elif len(valid_images) == 2:
                    auto_prompt = "이 두 이미지를 자연스럽게 합성해주세요. 두 이미지의 요소를 조화롭게 통합하여 하나의 이미지로 만들어주세요."
                else:
                    auto_prompt = "이 세 이미지를 창의적으로 합성해주세요. 모든 이미지의 주요 요소를 포함하되 자연스럽고 일관된 하나의 장면으로 만들어주세요."
            else:
                auto_prompt = preprocess_prompt(prompt, image1, image2, image3)
            
            # 이미지 생성 함수 호출
            result_img, status = process_images_with_prompt(image1, image2, image3, prompt)
            
            return result_img, status, auto_prompt
        except Exception as e:
            logger.exception("처리 중 오류 발생:")
            return None, f"오류 발생: {str(e)}", prompt
    
    submit_btn.click(
        fn=process_and_show_prompt,
        inputs=[image1_input, image2_input, image3_input, prompt_input],
        outputs=[output_image, output_text, prompt_display],
    )

    gr.Markdown(
        """
        ### 사용 방법:
        
        1. **자동 합성**: 이미지만 업로드하고 프롬프트를 비워두면 자동으로 합성됩니다
        2. **텍스트만으로 생성**: 이미지 없이 프롬프트만 입력해도 이미지가 생성됩니다
        3. **기능 사용**: 드롭다운에서 원하는 기능을 선택하면 자동으로 프롬프트가 입력됩니다
        4. **이미지 참조**: #1, #2, #3으로 각 이미지를 참조할 수 있습니다
        5. **일부 이미지만**: 필요한 이미지만 업로드해도 기능 실행이 가능합니다
        6. **재시도 기능**: 이미지 생성 중 오류 발생 시 자동으로 재시도합니다
        
        > **팁**: 기능을 선택한 후에도 프롬프트를 직접 수정할 수 있습니다
        """
    )

# 애플리케이션 실행
if __name__ == "__main__":
    demo.launch(share=True)