# Gradio UI 부분 수정
with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.HTML("""
        <div class="main-title">
            <h1>🎨 Advanced Image Object Extractor</h1>
            <p>Extract objects from images using text prompts or bounding boxes</p>
        </div>
    """)

    with gr.Tabs() as tabs:
        # Text 탭
        with gr.Tab("Extract by Text", id="tab_prompt"):
            with gr.Row(equal_height=True):
                with gr.Column(scale=1, min_width=400):
                    gr.HTML("<h3>Input Section</h3>")
                    input_image = gr.Image(
                        type="pil",
                        label="Upload Image",
                        interactive=True
                    )
                    with gr.Group():
                        text_prompt = gr.Textbox(
                            label="Object to Extract",
                            placeholder="Enter what you want to extract...",
                            interactive=True
                        )
                        background_prompt = gr.Textbox(
                            label="Background Generation Prompt (optional)",
                            placeholder="Describe the background you want...",
                            interactive=True
                        )
                    process_btn = gr.Button(
                        "Process Image",
                        variant="primary"
                    )

                with gr.Column(scale=1, min_width=400):
                    gr.HTML("<h3>Output Section</h3>")
                    output_image = ImageSlider(
                        label="Results Preview",
                        show_download_button=False
                    )
                    download_btn = gr.DownloadButton(
                        "Download Result"
                    )

            # Text 탭 Examples
            with gr.Accordion("Examples", open=False):
                text_examples = [
                    ["examples/text.jpg", "text", "white background"],
                    ["examples/black-lamp.jpg", "black lamp", "minimalist interior"]
                ]
                gr.Examples(
                    examples=text_examples,
                    inputs=[input_image, text_prompt, background_prompt],
                    outputs=[output_image, download_btn],
                    fn=process_prompt,
                    cache_examples=True
                )

        # Bounding Box 탭
        with gr.Tab("Extract by Box", id="tab_bb"):
            with gr.Row(equal_height=True):
                with gr.Column(scale=1, min_width=400):
                    gr.HTML("<h3>Input Section</h3>")
                    box_annotator = image_annotator(
                        image_type="pil",
                        disable_edit_boxes=False,  # 편집 가능하도록 변경
                        show_download_button=False,
                        show_share_button=False,
                        single_box=True,
                        label="Draw Box Around Object",
                        interactive=True
                    )
                    box_process_btn = gr.Button(
                        "Extract Selection",
                        variant="primary"
                    )

                with gr.Column(scale=1, min_width=400):
                    gr.HTML("<h3>Output Section</h3>")
                    box_output_image = ImageSlider(
                        label="Results Preview",
                        show_download_button=False
                    )
                    box_download_btn = gr.DownloadButton(
                        "Download Result"
                    )

            # Bounding Box Examples
            with gr.Accordion("Examples", open=False):
                box_examples = [
                    {
                        "image": "examples/text.jpg",
                        "boxes": [{"xmin": 51, "ymin": 511, "xmax": 639, "ymax": 1255}]
                    },
                    {
                        "image": "examples/black-lamp.jpg",
                        "boxes": [{"xmin": 88, "ymin": 148, "xmax": 700, "ymax": 1414}]
                    }
                ]
                gr.Examples(
                    examples=box_examples,
                    inputs=[box_annotator],
                    outputs=[box_output_image, box_download_btn],
                    fn=process_bbox,
                    cache_examples=True
                )

    # Event handlers
    def update_button_state(img, prompt):
        return gr.Button.update(interactive=bool(img and prompt))

    # Text 탭 이벤트
    input_image.change(
        fn=update_button_state,
        inputs=[input_image, text_prompt],
        outputs=[process_btn]
    )
    text_prompt.change(
        fn=update_button_state,
        inputs=[input_image, text_prompt],
        outputs=[process_btn]
    )
    process_btn.click(
        fn=process_prompt,
        inputs=[input_image, text_prompt, background_prompt],
        outputs=[output_image, download_btn]
    )

    # Bounding Box 탭 이벤트
    box_annotator.change(
        fn=lambda x: gr.Button.update(interactive=bool(x)),
        inputs=[box_annotator],
        outputs=[box_process_btn]
    )
    box_process_btn.click(
        fn=process_bbox,
        inputs=[box_annotator],
        outputs=[box_output_image, box_download_btn]
    )

# CSS 스타일 업데이트
css = """
footer {display: none}
.main-title {
    text-align: center;
    margin: 2em 0;
    padding: 1em;
    background: #f7f7f7;
    border-radius: 10px;
}
.main-title h1 {
    color: #2196F3;
    font-size: 2.5em;
    margin-bottom: 0.5em;
}
.main-title p {
    color: #666;
    font-size: 1.2em;
}
.container {
    max-width: 1200px;
    margin: auto;
    padding: 20px;
}
.tabs {
    margin-top: 1em;
}
.input-group {
    background: white;
    padding: 1em;
    border-radius: 8px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.output-group {
    background: white;
    padding: 1em;
    border-radius: 8px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
button.primary {
    background: #2196F3;
    border: none;
    color: white;
    padding: 0.5em 1em;
    border-radius: 4px;
    cursor: pointer;
    transition: background 0.3s ease;
}
button.primary:hover {
    background: #1976D2;
}
"""

# Launch settings
demo.queue(max_size=30, api_open=False)
demo.launch(
    show_api=False,
    share=False,
    server_name="0.0.0.0",
    server_port=7860,
    show_error=True
)