import json
import os
import time
import uuid
import tempfile
from PIL import Image
import gradio as gr
import base64
import mimetypes
import logging

from google import genai
from google.genai import types

# .env 파일에 저장된 환경변수 로드 (python-dotenv 설치 필요: pip install python-dotenv)
from dotenv import load_dotenv
load_dotenv()

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


def save_binary_file(file_name, data):
    logger.debug(f"파일에 이진 데이터 저장 중: {file_name}")
    with open(file_name, "wb") as f:
        f.write(data)
    logger.debug(f"파일 저장 완료: {file_name}")


def generate(text, file_name, model="gemini-2.0-flash-exp-image-generation"):
    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', 파일명: '{file_name}', 모델: '{model}'")

    try:
        # API 키는 환경변수에서 불러옴
        effective_api_key = os.environ.get("GEMINI_API_KEY")
        if effective_api_key:
            logger.debug("환경변수에서 API 키 불러옴")
        else:
            logger.error("API 키가 환경변수에 설정되지 않았습니다.")
            raise ValueError("API 키가 필요합니다.")

        client = genai.Client(api_key=effective_api_key)
        logger.debug("Gemini 클라이언트 초기화 완료.")

        # 파일 업로드
        files = [
            client.files.upload(file=file_name),
        ]
        logger.debug(f"파일 업로드 완료. URI: {files[0].uri}, MIME 타입: {files[0].mime_type}")

        # 컨텐츠 객체 생성: 파일 URI와 텍스트 프롬프트를 함께 포함
        contents = [
            types.Content(
                role="user",
                parts=[
                    types.Part.from_uri(
                        file_uri=files[0].uri,
                        mime_type=files[0].mime_type,
                    ),
                    types.Part.from_text(text=text),
                ],
            ),
        ]
        logger.debug(f"컨텐츠 객체 생성 완료: {contents}")

        generate_content_config = types.GenerateContentConfig(
            temperature=1,
            top_p=0.95,
            top_k=40,
            max_output_tokens=8192,
            response_modalities=[
                "image",
                "text",
            ],
            response_mime_type="text/plain",
        )
        logger.debug(f"생성 설정: {generate_content_config}")

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            temp_path = tmp.name
            logger.debug(f"임시 파일 생성됨: {temp_path}")

            response_stream = client.models.generate_content_stream(
                model=model,
                contents=contents,
                config=generate_content_config,
            )

            logger.debug("응답 스트림 처리 시작...")
            for chunk in response_stream:
                if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
                    logger.warning("chunk에 후보, 컨텐츠, 또는 파트가 없습니다. 건너뜁니다.")
                    continue

                inline_data = chunk.candidates[0].content.parts[0].inline_data
                if inline_data:
                    save_binary_file(temp_path, inline_data.data)
                    logger.info(f"MIME 타입 {inline_data.mime_type}의 파일이 저장됨: {temp_path} (프롬프트: {text})")
                else:
                    logger.info(f"수신된 텍스트: {chunk.text}")
                    print(chunk.text)

                logger.debug(f"Raw chunk: {chunk}")

        del files
        logger.debug("업로드된 파일 정보 삭제 완료.")
        return temp_path

    except Exception as e:
        logger.exception("이미지 생성 중 오류 발생:")
        return None  # 오류 발생 시 None 반환


def process_image_and_prompt(composite_pil, prompt):
    logger.debug(f"process_image_and_prompt 함수 시작 - 프롬프트: '{prompt}'")
    try:
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            composite_path = tmp.name
            composite_pil.save(composite_path)
            logger.debug(f"합성 이미지 저장 완료: {composite_path}")

        file_name = composite_path
        input_text = prompt
        model = "gemini-2.0-flash-exp-image-generation"

        gemma_edited_image_path = generate(text=input_text, file_name=file_name, model=model)

        if gemma_edited_image_path:
            logger.debug(f"이미지 생성 완료. 경로: {gemma_edited_image_path}")
            result_img = Image.open(gemma_edited_image_path)
            if result_img.mode == "RGBA":
                result_img = result_img.convert("RGB")
            return [result_img]
        else:
            logger.error("generate 함수에서 None 반환됨.")
            return []  # 오류 시 빈 리스트 반환

    except Exception as e:
        logger.exception("process_image_and_prompt 함수에서 오류 발생:")
        return []  # 오류 시 빈 리스트 반환


# --- Gradio 인터페이스 구성 ---
with gr.Blocks() as demo:
    gr.HTML(
        """
        <div style='display: flex; align-items: center; justify-content: center; gap: 20px'>
            <div style="background-color: var(--block-background-fill); border-radius: 8px">
                <img src="https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png" style="width: 100px; height: 100px;">
            </div>
            <div>
                <h1>Gemini를 이용한 이미지 편집</h1>
                <p>Gemini API 키는 환경변수(GEMINI_API_KEY)로 설정되어 있습니다.</p>
            </div>
        </div>
        """
    )
    gr.Markdown("이미지를 업로드하고, 편집할 내용을 입력하세요.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="이미지 업로드", image_mode="RGBA")
            prompt_input = gr.Textbox(
                lines=2,
                placeholder="편집할 내용을 입력하세요...",
                label="편집 프롬프트"
            )
            submit_btn = gr.Button("이미지 편집 실행")
        with gr.Column():
            output_gallery = gr.Gallery(label="편집 결과")

    submit_btn.click(
        fn=process_image_and_prompt,
        inputs=[image_input, prompt_input],
        outputs=output_gallery,
    )

# --- 테스트 코드 ---
# 테스트용 더미 이미지 (실제 이미지로 대체 가능)
dummy_image = Image.new("RGBA", (100, 100), color="red")
dummy_prompt = "이미지를 파란색으로 변경해줘"

logger.info("process_image_and_prompt 함수를 직접 호출합니다...")
result = process_image_and_prompt(dummy_image, dummy_prompt)

if result:
    logger.info(f"직접 호출 성공. 결과: {result}")
else:
    logger.error("직접 호출 실패.")

demo.launch(share=True)