import gradio as gr
import replicate
from PIL import Image
import tempfile
import uuid
import requests
from io import BytesIO
import os

# 환경 변수에서 API 토큰, 비밀번호, 모델 이름들 가져오기
REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
PASSWORD = os.getenv("PASSWORD")
UPSCALER_MODEL = os.getenv("UPSCALER_MODEL")
TEXT_REMOVAL_MODEL = os.getenv("TEXT_REMOVAL_MODEL")

def upscale_image(image_path, output_format):
    """Clarity Upscaler를 사용하여 이미지 업스케일링 (모든 매개변수 고정)"""
    try:
        client = replicate.Client(api_token=REPLICATE_API_TOKEN)
        
        with open(image_path, "rb") as file:
            input_data = {
                "image": file,
                "scale_factor": 2,  # 고정: 2배 확대
                "resemblance": 0.8,  # 고정: 높은 원본 유사도
                "creativity": 0.2,   # 고정: 낮은 창의성
                "output_format": output_format.lower(),
                "prompt": "high quality, detailed, sharp",
                "negative_prompt": "(worst quality, low quality, normal quality:2)"
            }
            
            output = client.run(
                UPSCALER_MODEL,
                input=input_data
            )
        
        # 결과 이미지 다운로드
        if output and isinstance(output, list) and len(output) > 0:
            response = requests.get(output[0])
            if response.status_code == 200:
                result_image = Image.open(BytesIO(response.content))
                
                # 임시 파일로 저장
                ext = output_format.lower()
                upscaled_filename = f"upscaled_temp_{uuid.uuid4()}.{ext}"
                
                if ext == 'jpg' and result_image.mode == 'RGBA':
                    result_image = result_image.convert('RGB')
                
                result_image.save(upscaled_filename, format='JPEG' if ext=='jpg' else 'PNG', quality=95 if ext=='jpg' else None)
                
                return upscaled_filename
            else:
                return None
        else:
            return None
        
    except Exception as e:
        return None

def remove_text_from_image(input_image, password, output_format, upscale_option):
    if not password:
        return None, "비밀번호를 입력하세요.", None
    
    if password != PASSWORD:
        return None, "비밀번호가 올바르지 않습니다.", None
    
    if input_image is None:
        return None, "이미지를 업로드하세요.", None
    
    try:
        # Replicate 클라이언트 설정
        client = replicate.Client(api_token=REPLICATE_API_TOKEN)
        
        log_message = f"이미지 처리 시작...\n"
        
        # Option 2: 로컬 파일 방식 사용
        with open(input_image, "rb") as file:
            input_file = file
            
            log_message += f"파일 읽기 완료\n"
            
            # Replicate API 호출 - output_format 매개변수 추가
            input_data = {
                "input_image": input_file,
                "output_format": output_format  # JPG 또는 PNG
            }
            
            log_message += f"텍스트 제거 API 호출 중... (출력 포맷: {output_format})\n"
            
            output = client.run(
                TEXT_REMOVAL_MODEL,
                input=input_data
            )
        
        log_message += f"텍스트 제거 API 응답 수신 완료\n"
        
        # 결과 이미지 처리
        temp_filename = f"temp_output.{output_format.lower()}"
        with open(temp_filename, "wb") as file:
            file.write(output.read())
        
        # 결과 이미지 로드
        result_image = Image.open(temp_filename)
        
        # 화질 개선 적용 (선택사항)
        final_image = result_image
        if upscale_option == "적용":
            log_message += f"Clarity Upscaler로 화질 개선 시작... (2배 확대, 유사도: 0.8, 창의성: 0.2)\n"
            upscaled_path = upscale_image(temp_filename, output_format)
            
            if upscaled_path:
                final_image = Image.open(upscaled_path)
                log_message += f"화질 개선 완료\n"
            else:
                log_message += f"화질 개선 실패, 원본 이미지 사용\n"
        
        # 다운로드용 파일명
        download_filename = f"text_removed.{output_format.lower()}"
        final_image.save(download_filename)
        
        log_message += f"텍스트 제거 완료! (API에서 {output_format} 포맷으로 처리됨)"
        
        return final_image, log_message, download_filename
        
    except Exception as e:
        error_message = f"오류 발생: {str(e)}"
        return None, error_message, None

# Gradio 인터페이스
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="filepath", label="입력 이미지")
            api_password = gr.Textbox(
                label="비밀번호", 
                type="password"
            )
            output_format = gr.Radio(
                choices=["jpg", "png"],
                value="jpg",
                label="출력 포맷"
            )
            upscale_option = gr.Radio(
                choices=["없음", "적용"],
                value="없음",
                label="화질 개선 (2배 확대)"
            )
            submit_btn = gr.Button("텍스트 제거", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="처리된 이미지")
            log_output = gr.Textbox(label="로그", lines=10, max_lines=15)
            download_file = gr.File(label="다운로드")
    
    submit_btn.click(
        fn=remove_text_from_image,
        inputs=[input_image, api_password, output_format, upscale_option],
        outputs=[output_image, log_output, download_file]
    )

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