import gradio as gr
import replicate
import os
import requests
import tempfile
import logging
import base64

logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def process_image(password, input_image):
    # 환경변수에서 비밀번호 가져오기
    correct_password = os.getenv("APP_PASSWORD")
    if not correct_password:
        logger.error("APP_PASSWORD 환경변수가 설정되지 않았습니다.")
        raise ValueError("서버 설정 오류입니다.")
    
    # 비밀번호 검증
    if password != correct_password:
        raise ValueError("잘못된 비밀번호입니다.")
    
    # Replicate API 토큰 확인
    if not os.getenv("REPLICATE_API_TOKEN"):
        logger.error("REPLICATE_API_TOKEN이 설정되지 않았습니다.")
        return None, None
    
    # 환경변수에서 모델 이름 가져오기
    model_name = os.getenv("REPLICATE_MODEL")
    if not model_name:
        logger.error("REPLICATE_MODEL 환경변수가 설정되지 않았습니다.")
        return None, None
    
    if input_image is None:
        logger.error("입력 이미지가 없습니다.")
        return None, None
    
    try:
        # base64로 변환
        with open(input_image, "rb") as f:
            data = base64.b64encode(f.read()).decode()
        
        image_uri = f"data:image/png;base64,{data}"
        
        # 환경변수에서 가져온 모델로 실행
        output = replicate.run(
            model_name,
            input={"image": image_uri}
        )
        
        # 결과 저장
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
            if hasattr(output, "read"):
                tmp.write(output.read())
            elif isinstance(output, (list, tuple)) and output:
                resp = requests.get(output[0])
                resp.raise_for_status()
                tmp.write(resp.content)
            
            out_path = tmp.name
        
        return out_path, out_path
        
    except replicate.exceptions.ReplicateError as re:
        logger.error(f"API 오류: {re}")
        return None, None
    except Exception as e:
        logger.error(f"예상치 못한 오류: {e}")
        return None, None

# Gradio 인터페이스 생성
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    with gr.Row():
        with gr.Column():
            password_input = gr.Textbox(
                label="비밀번호",
                type="password",
                placeholder="비밀번호를 입력하세요"
            )
            input_image = gr.Image(
                type="filepath",
                label="입력 이미지",
                interactive=True
            )
            process_btn = gr.Button("실행", variant="primary")
            
        with gr.Column():
            output_image = gr.Image(
                type="filepath",
                label="결과 이미지",
                interactive=False
            )
            download_btn = gr.DownloadButton(
                label="이미지 다운로드"
            )
    
    # click 시 (표시용, 다운로드용) 두 개 리턴
    process_btn.click(
        fn=process_image,
        inputs=[password_input, input_image],
        outputs=[output_image, download_btn]
    )
    
    gr.Markdown("---")

if __name__ == "__main__":
    demo.launch()