import gradio as gr
import replicate
import os
import requests
from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
import io
import tempfile
import numpy as np

# Replicate API 토큰 설정
REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
if REPLICATE_API_TOKEN:
    os.environ["REPLICATE_API_TOKEN"] = REPLICATE_API_TOKEN

def upload_image_to_temp_url(image):
    """이미지를 임시 파일로 저장"""
    if image is None:
        return None
    
    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp_file:
        image.save(tmp_file.name, format='PNG')
        return tmp_file.name

def remove_background_with_ai(image_path):
    """AI 배경 제거"""
    if not REPLICATE_API_TOKEN:
        return None, "Replicate API 토큰이 필요합니다."
    
    try:
        output = replicate.run(
            "851-labs/background-remover:a029dff38972b5fda4ec5d75d7d1cd25aeff621d2cf4946a41055d7db66b80bc",
            input={"image": open(image_path, "rb")}
        )
        
        response = requests.get(output)
        if response.status_code == 200:
            return Image.open(io.BytesIO(response.content)), "배경 제거 완료"
        else:
            return None, "배경 제거 실패"
            
    except Exception as e:
        return None, f"배경 제거 오류: {str(e)}"

def upscale_with_clarity(image_path):
    """화질 개선 업스케일러"""
    if not REPLICATE_API_TOKEN:
        return None, "Replicate API 토큰이 필요합니다."
    
    try:
        output = replicate.run(
            "philz1337x/clarity-upscaler:dfad41707589d68ecdccd1dfa600d55a208f9310748e44bfe35b4a6291453d5e",
            input={"image": open(image_path, "rb")}
        )
        
        # output은 리스트이므로 첫 번째 항목 사용
        response = requests.get(output[0])
        if response.status_code == 200:
            return Image.open(io.BytesIO(response.content)), "화질 개선 완료"
        else:
            return None, "화질 개선 실패"
            
    except Exception as e:
        return None, f"화질 개선 오류: {str(e)}"

def change_product_angle(image_path):
    """카메라 위치를 바꿔서 다른 앵글로 촬영"""
    if not REPLICATE_API_TOKEN:
        return None, "Replicate API 토큰이 필요합니다."
    
    try:
        # Gen4 Image로 카메라 위치 변경 (파일 직접 업로드)
        angle_prompt = "Professional product photography of @product from 45-degree side angle, three-quarter view, diagonal perspective, different camera position, clean background"
        
        with open(image_path, "rb") as image_file:
            output = replicate.run(
                "runwayml/gen4-image",
                input={
                    "prompt": angle_prompt,
                    "aspect_ratio": "1:1",
                    "resolution": "1080p",
                    "reference_images": [image_file],
                    "reference_tags": ["product"]
                }
            )
        
        response = requests.get(output)
        if response.status_code == 200:
            return Image.open(io.BytesIO(response.content)), "카메라 앵글 변경 완료 (Gen4)"
        else:
            return None, "카메라 앵글 변경 실패"
            
    except Exception as e:
        return None, f"카메라 앵글 변경 오류: {str(e)}"

def add_shadow_only(image_path):
    """배경 제거된 상품에 그림자만 추가 (완전한 흰색 배경)"""
    if not REPLICATE_API_TOKEN:
        return None, "Replicate API 토큰이 필요합니다."
    
    try:
        # Gen4 Image로 그림자 추가 (파일 직접 업로드)
        shadow_prompt = "Professional product photography of @product on pure white background with natural soft drop shadow, clean white backdrop, subtle shadow effect, e-commerce style"
        
        with open(image_path, "rb") as image_file:
            output = replicate.run(
                "runwayml/gen4-image",
                input={
                    "prompt": shadow_prompt,
                    "aspect_ratio": "1:1", 
                    "resolution": "1080p",
                    "reference_images": [image_file],
                    "reference_tags": ["product"]
                }
            )
        
        response = requests.get(output)
        if response.status_code == 200:
            return Image.open(io.BytesIO(response.content)), "그림자 추가 완료 (Gen4)"
        else:
            return None, "그림자 추가 실패"
            
    except Exception as e:
        return None, f"그림자 추가 오류: {str(e)}"

def process_image(image, remove_bg, upscale_quality):
    """메인 처리 함수"""
    if image is None:
        return None, "이미지를 업로드해주세요."
    
    try:
        result_image = image.copy()
        status_messages = []
        
        # 1단계: 상품 각도 변경
        temp_path = upload_image_to_temp_url(result_image)
        if temp_path:
            angled_img, angle_msg = change_product_angle(temp_path)
            if angled_img:
                result_image = angled_img
                status_messages.append(f"✅ {angle_msg}")
            else:
                status_messages.append(f"❌ {angle_msg}")
            os.unlink(temp_path)
        
        # 2단계: 배경 제거 (누끼)
        if remove_bg:
            temp_path = upload_image_to_temp_url(result_image)
            if temp_path:
                bg_removed_img, bg_msg = remove_background_with_ai(temp_path)
                if bg_removed_img:
                    result_image = bg_removed_img
                    status_messages.append(f"✅ {bg_msg}")
                else:
                    status_messages.append(f"❌ {bg_msg}")
                os.unlink(temp_path)
        
        # 3단계: 그림자만 추가 (완전한 흰색 배경)
        temp_path = upload_image_to_temp_url(result_image)
        if temp_path:
            shadow_img, shadow_msg = add_shadow_only(temp_path)
            if shadow_img:
                result_image = shadow_img
                status_messages.append(f"✅ {shadow_msg}")
            else:
                status_messages.append(f"❌ {shadow_msg}")
            os.unlink(temp_path)
        
        # 4단계: 화질 개선 (이미지 생성 후)
        if upscale_quality:
            temp_path = upload_image_to_temp_url(result_image)
            if temp_path:
                upscaled_img, upscale_msg = upscale_with_clarity(temp_path)
                if upscaled_img:
                    result_image = upscaled_img
                    status_messages.append(f"✅ {upscale_msg}")
                else:
                    status_messages.append(f"❌ {upscale_msg}")
                os.unlink(temp_path)
        
        # 5단계: 최종 크기 조정
        result_image = result_image.resize((1000, 1000), Image.Resampling.LANCZOS)
        status_messages.append("✅ 1000x1000 크기 조정 완료")
        
        return result_image, "\n".join(status_messages)
        
    except Exception as e:
        return None, f"처리 중 오류: {str(e)}"

# Gradio 인터페이스
def create_interface():
    with gr.Blocks(title="쿠팡 썸네일 생성기", theme=gr.themes.Soft()) as iface:
        
        with gr.Row():
            with gr.Column(scale=1):
                # 이미지 업로드
                input_image = gr.Image(
                    label="상품 이미지 업로드", 
                    type="pil",
                    height=400
                )
                
                # 선택 기능 2가지
                with gr.Group():
                    remove_bg = gr.Checkbox(
                        label="배경 제거", 
                        value=True,
                        info="AI로 배경을 자동 제거합니다"
                    )
                    upscale_quality = gr.Checkbox(
                        label="화질 개선기", 
                        value=True,
                        info="AI로 화질을 향상시킵니다"
                    )
                
                # 실행 버튼
                process_btn = gr.Button("🚀 썸네일 생성", variant="primary", size="lg")
            
            with gr.Column(scale=1):
                # 출력 이미지
                output_image = gr.Image(
                    label="쿠팡 썸네일 결과", 
                    height=400
                )
                status_output = gr.Textbox(
                    label="처리 상태", 
                    lines=6,
                    max_lines=10
                )
        
        # 버튼 클릭 이벤트
        process_btn.click(
            process_image,
            inputs=[input_image, remove_bg, upscale_quality],
            outputs=[output_image, status_output]
        )
    
    return iface

# 앱 실행
if __name__ == "__main__":
    app = create_interface()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True
    )