src/ui/components.py

"""
UI组件模块
"""

import gradio as gr
from typing import List, Tuple

from ..config import (
    UI_CONFIG, 
    EXAMPLE_PROMPTS, 
    ASPECT_RATIOS
)
from ..utils import handle_file_upload, create_image_html


class UIComponents:
    """UI组件类"""
    
    def __init__(self):
        self.delete_buttons = []
    
    def create_header(self) -> gr.HTML:
        """创建页面头部"""
        return gr.HTML(f"""
            <h1 class="logo-text">{UI_CONFIG['APP_TITLE']}</h1>
            <p class="subtitle">{UI_CONFIG['APP_SUBTITLE']}</p>
            <div class="mode-indicator">
                {UI_CONFIG['APP_DESCRIPTION']}
            </div>
        """)
    
    def create_info_box(self) -> gr.HTML:
        """创建信息提示框"""
        return gr.HTML("""
            <div class="info-box">
                <strong>Usage Instructions:</strong><br>
                • Get your Veo 3 API Key   <a href="https://kie.ai/playground/veo3" target="_blank">👉 here 👈</a><br>
                • Select your model: Choose Veo 3 Fast for more cost-effective video generation (Free test) or Veo 3 Quality for higher-quality videos.<br>
                • Enter your prompt (text-to-video) or upload your image (image-to-video).<br>
                • Click Generate and wait ~1–2 minutes for processing
            </div>
        """)
    
    def create_input_components(self) -> Tuple[gr.Textbox, gr.Textbox, gr.HTML, gr.File, List[gr.Button], gr.Dropdown, gr.Number, gr.Button]:
        """创建输入组件"""
        # API Key输入
        api_key = gr.Textbox(
            label="API Key",
            placeholder="Please enter your KIE AI API Key",
            type="password",
            elem_classes="api-key-input"
        )

        # 提示词输入
        prompt = gr.Textbox(
            label="Video Prompt",
            placeholder="Describe the video you want to generate, e.g.: A dog playing in a park...",
            lines=3,
            value=UI_CONFIG["DEFAULT_PROMPT"],
            elem_classes="prompt-input"
        )
        
        # 图片上传区域 - 放在提示词和宽高比之间
        with gr.Group(elem_classes="image-upload-container"):
            gr.Markdown("### 📸 Image Upload (Optional)")
            gr.Markdown("*Upload 1 image for image-to-video generation, or leave empty for text-to-video*")

            # 自定义图片展示区
            image_display = gr.HTML(
                value="<div id='image-display-area'><div class='no-images'>No image uploaded - will generate from text only</div></div>",
                elem_id="image-display"
            )

            # 上传按钮
            file_upload = gr.File(
                show_label=False,
                file_count="single",
                file_types=["image"],
                type="filepath",
                height=120,
            )

            # 删除按钮组
            with gr.Row(elem_id="delete-buttons-row"):
                delete_label = gr.Markdown("**Delete Images:**", visible=True, elem_id="delete-label")
                delete_buttons = []
                for i in range(1):  # 最多支持1张图片
                    btn = gr.Button(f"Delete {i + 1}", visible=True, size="sm", elem_id=f"delete-btn-{i}")
                    delete_buttons.append(btn)
        
        # 宽高比选择
        aspect_ratio = gr.Dropdown(
            choices=ASPECT_RATIOS,
            value=UI_CONFIG["DEFAULT_ASPECT_RATIO"],
            label="Aspect Ratio",
            info="16:9 for landscape, 9:16 for portrait"
        )
        
        # 种子输入区域 - 使用Row让按钮在输入框右边
        with gr.Row():
            seeds = gr.Number(
                label="Seed (Optional)",
                value=10001,
                info="Random seed for reproducible results (10000-99999)",
                scale=4
            )
            random_seed_btn = gr.Button(
                "🎲",
                size="sm",
                scale=1,
                elem_id="random-seed-btn"
            )
        
        # 添加一个状态来保存上传的文件路径
        uploaded_file_state = gr.State(None)
        
        return api_key, prompt, image_display, file_upload, delete_buttons, aspect_ratio, seeds, random_seed_btn, uploaded_file_state

    def create_output_section(self) -> Tuple[gr.Video, gr.Textbox]:
        """创建输出区域"""
        # 输出视频
        output_video = gr.Video(
            show_label=False,
            elem_id="output-video",
            height=400,
            container=True
        )

        # 状态信息
        status = gr.Textbox(
            label="Processing Status",
            interactive=False,
            lines=2,
            value="Ready, please enter a prompt to generate video..."
        )
        
        return output_video, status
    
    def create_examples(self) -> gr.Examples:
        """创建示例"""
        return gr.Examples(
            examples=[[prompt, None] for prompt in EXAMPLE_PROMPTS],
            inputs=[],  # 将在主应用中设置
            label="Video Prompt Examples"
        )
    
    def create_generate_button(self) -> gr.Button:
        """创建生成按钮"""
        return gr.Button(
            "🚀 Start Generation",
            variant="primary",
            size="lg"
        )
    
    def setup_file_upload_handlers(self, file_upload, image_display, delete_buttons, uploaded_file_state):
        """设置文件上传处理器"""
        def on_file_upload(new_files):
            if not new_files:
                return gr.update(value=None), "<div id='image-display-area'><div class='no-images'>No image uploaded - will generate from text only</div></div>", gr.update(visible=False), None
            
            # 只处理第一张图片
            file_path = new_files[0] if isinstance(new_files, list) else new_files
            
            # 显示图片预览，实际上传在生成时进行
            html_content = create_image_html([file_path])
            return gr.update(value=None), html_content, gr.update(visible=True, value="Delete Image 1"), file_path

        file_upload.upload(
            fn=on_file_upload,
            inputs=[file_upload],
            outputs=[file_upload, image_display, delete_buttons[0], uploaded_file_state]
        )
    
    def setup_delete_handlers(self, delete_buttons, image_display, uploaded_file_state):
        """设置删除按钮处理器"""
        def delete_image():
            return "<div id='image-display-area'><div class='no-images'>No image uploaded - will generate from text only</div></div>", gr.update(visible=False), None

        # 绑定删除按钮事件
        delete_buttons[0].click(
            fn=delete_image,
            outputs=[image_display, delete_buttons[0], uploaded_file_state]
        )