import gradio as gr
from infra import image_processing_single, image_processing_batch

with gr.Blocks(theme='freddyaboulton/dracula_revamped', title='Image Generator') 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():
                product_image_size = gr.Slider(512, 1024, value=800, label="제품 이미지 사이즈")
                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"):
            with gr.Column():
                product_image_size_batch = gr.Slider(512, 1024, value=800, label="제품 이미지 사이즈")
                gr.HTML("<div>링크1(엔터) 링크2(엔터) ... 으로 구성된 .txt 파일을 업로드</div>")
                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, product_image_size], outputs=[preview, output_zip], api_name="image_processing_single")
    btn_batch.click(fn=image_processing_batch, inputs=[background_image, links, product_image_size_batch], outputs=[preview, output_zip], api_name="image_processing_batch")

demo.launch()