import requests
import gradio as gr
import zipfile
import os
import tempfile
from PIL import Image
from io import BytesIO
from rembg import remove

def download_image(image_url: str) -> Image.Image:
    # 이미지 다운로드
    response = requests.get(image_url)
    
    if response.status_code == 200:
        original_image = Image.open(BytesIO(response.content))
        return original_image
    else:
        raise Exception(f"Failed to download image. Status code: {response.status_code}")

def remove_background(image: Image.Image) -> Image.Image:
    # 이미지 누끼따기
    try:
        removebg_image = remove(
            image,
            post_process_mask=True,
            alpha_matting=True, 
            alpha_matting_foreground_threshold=270,
            alpha_matting_background_threshold=30, 
            alpha_matting_erode_size=15)
        return removebg_image
    
    except Exception as e:
        print(f"Failed to remove background: {e}")
        return None
    
def crop_image(image: Image.Image) -> Image.Image:
    # 이미지 크롭
    try:
        # 알파 채널을 사용하여 이미지의 경계 영역 찾기
        bbox = image.getbbox()
        if bbox:
            # 경계 상자로 이미지 크롭
            cropped_image = image.crop(bbox)
            return cropped_image
        else:
            print("No bounding box found.")
            return image
    
    except Exception as e:
        print(f"Failed to crop image: {e}")
        return None
      
def resize_image(image: Image.Image, max_size: int) -> Image.Image:
    # 이미지 크기 조정
    try:
        # 이미지의 현재 너비와 높이 가져오기
        width, height = image.size
        
        # 너비와 높이 중 더 큰 쪽의 비율에 맞춰 크기를 조정
        if width > height:
            new_width = max_size
            new_height = int((max_size / width) * height)
        else:
            new_height = max_size
            new_width = int((max_size / height) * width)
        
        resized_image = image.resize((new_width, new_height))
        return resized_image
    
    except Exception as e:
        print(f"Failed to resize image: {e}")
        return None
    
def paste_to_background_type_a(background: Image.Image, product: Image.Image) -> Image.Image:
    # 배경에 제품 이미지 합성
    try:
        # 배경 이미지 크기 가져오기
        bg_width, bg_height = background.size
        
        # 제품 이미지 크기 가져오기
        product_width, product_height = product.size
        
        # 제품 이미지를 배경 이미지 중앙에 위치시키기
        offset = ((bg_width - product_width) // 2, (bg_height - product_height) // 2)
        
        # 알파 채널을 고려하여 합성
        background.paste(product, offset, mask=product)
        return background
    
    except Exception as e:
        print(f"Failed to paste product image to background: {e}")
        return None
    
def paste_to_background_type_b(background: Image.Image, product: Image.Image, margin: int = 10) -> Image.Image:
    try:
        # 배경 이미지 크기 가져오기
        bg_width, bg_height = background.size
        
        # 제품 이미지 크기 가져오기
        product_width, product_height = product.size
        
        # 두 제품 이미지를 위한 전체 너비 계산 (제품 이미지 두 개 + 마진)
        total_width = (product_width * 2) + margin
        
        if total_width > bg_width:
            raise ValueError("Background is too narrow to fit two product images with the specified margin.")
        
        # 첫 번째 제품 이미지의 왼쪽 상단 좌표 계산
        left_offset = (bg_width - total_width) // 2
        top_offset = (bg_height - product_height) // 2
        
        # 두 번째 제품 이미지의 왼쪽 상단 좌표 계산
        right_offset = left_offset + product_width + margin
        
        # 첫 번째 제품 이미지 배경에 합성
        background.paste(product, (left_offset, top_offset), mask=product)
        
        # 두 번째 제품 이미지 배경에 합성
        background.paste(product, (right_offset, top_offset), mask=product)
        return background
    
    except Exception as e:
        print(f"Failed to paste product images to background: {e}")
        return None

def paste_to_background_type_c(background: Image.Image, product: Image.Image, margin: int = 10, top_margin: int = 15) -> Image.Image:
    try:
        # 배경 이미지 크기 가져오기
        bg_width, bg_height = background.size
        
        # 제품 이미지 크기 가져오기
        product_width, product_height = product.size
        
        # 먼저 type_b 형태로 아래 두 제품 이미지를 배치
        background_with_two_products = paste_to_background_type_b(background, product, margin)

        # 중앙 상단에 위치할 제품 이미지의 오프셋 계산
        center_offset_x = (bg_width - product_width) // 2
        center_offset_y = ((bg_height - product_height) // 2) + top_margin
        
        if center_offset_y < 0:
            raise ValueError("Background is too small to fit three product images with the specified top margin.")
        
        # 세 번째 제품 이미지 배경에 합성 (중앙 상단)
        background_with_two_products.paste(product, (center_offset_x, center_offset_y), mask=product)
        
        # 결과 이미지 저장 및 반환
        return background_with_two_products
    
    except Exception as e:
        print(f"Failed to paste product images to background: {e}")
        return None
    
def create_zip_file(images: list) -> str:
    # 임시 파일 생성
    with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as temp_zip:
        with zipfile.ZipFile(temp_zip, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for i, image in enumerate(images):
                # 이미지를 메모리 버퍼에 저장
                image_buffer = BytesIO()
                image.save(image_buffer, format="PNG")
                image_buffer.seek(0)
                
                # 버퍼 내용을 ZIP 파일에 추가
                zipf.writestr(f"image_{i + 1}.png", image_buffer.getvalue())
        
        # 임시 파일의 경로를 반환
        temp_zip_path = temp_zip.name
    
    return temp_zip_path

def image_processing_single(background_image: Image.Image, product_image_url: str):
    # 각종 설정 값
    background_image_size = 1024 # 배경 이미지의 크기
    product_image_size = 650 # 배경이 제거된 제품 이미지의 크기
    type_b_margin = 15
    type_c_margin_1, type_c_margin_2 = 15, 45 # (1=가운데 두개의 마진, 2= 중앙 한개의 top 마진)
    
    # 배경 이미지 크기 조정
    background_image = background_image.resize((background_image_size, background_image_size))
    
    # 이미지 다운로드
    original_image = download_image(product_image_url)

    # 이미지 누끼따기
    removebg_image = remove_background(original_image)
    
    # 이미지 크롭
    cropped_image = crop_image(removebg_image)
    
    # 크롭된 이미지 원하는 사이즈로 resize
    resized_image = resize_image(cropped_image, product_image_size)
    
    # type_a 합성
    type_a_image = paste_to_background_type_a(background_image.copy(), resized_image)
    type_b_image = paste_to_background_type_b(background_image.copy(), resized_image, type_b_margin)
    type_c_image = paste_to_background_type_c(background_image.copy(), resized_image, type_c_margin_1, type_c_margin_2)
    
    # 결과 이미지 반환
    image_list = [type_a_image, type_b_image, type_c_image]
    zip_file_path = create_zip_file(image_list)
    
    return image_list, zip_file_path

def image_processing_batch(background_image: Image.Image, product_image_url_file_path: str):
    # 각종 설정 값
    product_image_size = 650 # 배경이 제거된 제품 이미지의 크기
    type_b_margin = 15
    type_c_margin_1, type_c_margin_2 = 15, 45 # (1=가운데 두개의 마진, 2= 중앙 한개의 top 마진)
    
    # file to url
    def file_to_list(file_path: str) -> list:
        # 파일을 열고 읽기
        with open(file_path, "r") as f:
            # 파일 내용을 읽어서 줄바꿈(엔터)로 구분된 각 URL을 리스트로 변환
            content = f.read()
            # 줄바꿈으로 URL 분리
            url_list = [url.strip() for url in content.splitlines() if url.strip()]

        return url_list
    
    product_image_url_list = file_to_list(product_image_url_file_path)
    
    preview_image_list, image_list = [], []
    for idx, product_image_url in enumerate(product_image_url_list):
        # 이미지 다운로드
        original_image = download_image(product_image_url)

        # 이미지 누끼따기
        removebg_image = remove_background(original_image)
        
        # 이미지 크롭
        cropped_image = crop_image(removebg_image)
        
        # 크롭된 이미지 원하는 사이즈로 resize
        resized_image = resize_image(cropped_image, product_image_size)
        
        # type_a 합성
        type_a_image = paste_to_background_type_a(background_image.copy(), resized_image)
        type_b_image = paste_to_background_type_b(background_image.copy(), resized_image, type_b_margin)
        type_c_image = paste_to_background_type_c(background_image.copy(), resized_image, type_c_margin_1, type_c_margin_2)
        
        image_list.append(type_a_image)
        image_list.append(type_b_image)
        image_list.append(type_c_image)
        if idx == 0:
            preview_image_list = [type_a_image, type_b_image, type_c_image]
            
    zip_file_path = create_zip_file(image_list)
    
    return preview_image_list, zip_file_path
        

with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.HTML("""
        <h1>Product Image Generator</h1>
        <p>1. 배경 이미지를 업로드해주세요.<p>
        <p>2-1. Single 모드의 경우, 이미지의 링크를 입력해주세요.<p>
        <p>2-2. Batch 모드의 경우, 이미지 링크들이 포함된 .txt 파일(형식 = 링크1(엔터) 링크2(엔터) ..)을 업로드해주세요.<p>
        <p>3. Generate Images 버튼을 눌러 제품 이미지를 생성해주세요.<p>
    """)
    
    with gr.Row(elem_classes="input_box"):
        with gr.Column():
            background_image = gr.Image(label="Background Image", type="pil", height=300)
            gr.Examples(
                examples=[["examples/background_1.png"], ["examples/background_2.png"], ["examples/background_3.png"], ["examples/background_4.png"]], # 예제 이미지 경로 리스트
                inputs=background_image,  # 예제를 클릭했을 때 해당 이미지가 이 input에 적용됨
                label="예제 배경 이미지")
        with gr.Tab("Single Mode"):
            with gr.Column():
                link = gr.Textbox(label="Image URL")
                gr.Examples(
                    examples=["https://url.kr/yx85be", "https://url.kr/if4wxo"],
                    inputs=link,
                    label="예제 이미지 링크"
                )
            btn_single = gr.Button("Generate Images")
        with gr.Tab("Batch Mode"):
            gr.HTML("<div>링크1(엔터) 링크2(엔터) ... 으로 구성된 .txt 파일을 업로드</div>")
            with gr.Column():
                links = gr.File(label="Image URLs (.txt file)", height=100)
                gr.Examples(
                    examples=[["examples/links.txt"]],
                    inputs=links,
                    label="예제 .txt 파일"
                )
            btn_batch = gr.Button("Generate Images")
                
    gr.HTML("""
        <br><br>
        <h1>생성 결과 확인</h1>
        <p>생성은 1개 이미지 당 15~30초 소요되며(배경제거 AI 모듈 처리시간), 결과는 다음과 같이 제공됩니다.</p>
        <p>1. 첫번째 이미지에 대한 프리뷰(3개 타입)</p> 
        <p>2. 결과 이미지 전체가 담긴 ZIP 파일(results.zip) 다운로드 링크가 제공됩니다</p>
    """)

    with gr.Row():
        preview = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", columns=[3], rows=[1], object_fit="contain", height="auto")
        output_zip = gr.File(label="Download Result Zip File")
        
    # evnets
    btn_single.click(fn=image_processing_single, inputs=[background_image, link], outputs=[preview, output_zip], api_name="image_processing_single")
    btn_batch.click(fn=image_processing_batch, inputs=[background_image, links], outputs=[preview, output_zip], api_name="image_processing_batch")

demo.launch()