app.py

import os
import asyncio
import threading
import logging
import tempfile
import gradio as gr
from pathlib import Path
from typing import List, Optional

from config import config
from processor import process_batch_sync
from api_clients import get_ai_client

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 确保输出目录存在
os.makedirs(config.OUTPUT_DIR, exist_ok=True)


def analyze_images(
    files,
    provider: str,
    api_key: str,
    base_url: str,
    model: str,
    custom_prompt: str,
    max_concurrent: int,
):
    """Gradio 主处理函数"""
    if not files:
        return None, "❌ 请上传至少一张图片"
    
    if not api_key:
        return None, "❌ 请输入 API Key"
    
    # 获取文件路径
    image_paths = []
    for f in files:
        if isinstance(f, str):
            image_paths.append(f)
        else:
            image_paths.append(f.name)
    
    try:
        zip_path, summary = process_batch_sync(
            image_paths=image_paths,
            provider=provider,
            api_key=api_key if api_key else None,
            base_url=base_url if base_url else None,
            model=model if model else None,
            custom_prompt=custom_prompt if custom_prompt.strip() else None,
            max_concurrent=int(max_concurrent),
        )
        return zip_path, summary
    
    except Exception as e:
        logger.error(f"处理失败: {e}")
        return None, f"❌ 处理失败: {str(e)}"


def create_gradio_app():
    """创建 Gradio 界面"""
    
    with gr.Blocks(
        title="🖼️ Image to Prompt",
        theme=gr.themes.Soft(),
        css="""
        .main-title { text-align: center; margin-bottom: 1em; }
        .result-box { font-family: monospace; }
        """
    ) as app:
        
        gr.Markdown(
            """
            # 🖼️ Image → Prompt 图片反推提示词
            上传图片，AI 自动生成 Stable Diffusion / Midjourney 提示词  
            支持批量处理，输出 ZIP 压缩包（图片 + 提示词文本）
            """,
            elem_classes="main-title"
        )
        
        with gr.Row():
            # ===== 左栏: 设置 =====
            with gr.Column(scale=1):
                gr.Markdown("### ⚙️ API 设置")
                
                provider = gr.Dropdown(
                    choices=["openai", "gemini", "claude", "qwen", "custom"],
                    value="openai",
                    label="AI 提供商",
                    info="选择视觉模型提供商"
                )
                
                api_key = gr.Textbox(
                    label="API Key",
                    type="password",
                    placeholder="sk-... / AIza...",
                    info="对应提供商的 API Key"
                )
                
                base_url = gr.Textbox(
                    label="API Base URL (可选)",
                    placeholder="https://api.openai.com/v1",
                    info="自定义 API 地址，仅 openai/qwen/custom 需要"
                )
                
                model = gr.Textbox(
                    label="模型名称 (可选)",
                    placeholder="gpt-4o",
                    info="留空使用默认模型"
                )
                
                max_concurrent = gr.Slider(
                    minimum=1, maximum=10, value=3, step=1,
                    label="并发数",
                    info="同时处理的图片数量"
                )
                
                gr.Markdown("### 📝 自定义提示词")
                custom_prompt = gr.Textbox(
                    label="提示词模板 (可选)",
                    lines=5,
                    placeholder="留空使用默认模板...\n\n例如: Analyze this image and generate a detailed Stable Diffusion prompt...",
                    info="自定义发送给 AI 的指令"
                )
            
            # ===== 右栏: 上传与结果 =====
            with gr.Column(scale=2):
                gr.Markdown("### 📤 上传图片")
                
                file_input = gr.Files(
                    label="拖拽或点击上传图片",
                    file_types=["image"],
                    file_count="multiple",
                )
                
                submit_btn = gr.Button(
                    "🚀 开始处理", 
                    variant="primary",
                    size="lg"
                )
                
                gr.Markdown("### 📊 处理结果")
                
                output_file = gr.File(
                    label="📦 下载 ZIP 压缩包",
                )
                
                output_summary = gr.Textbox(
                    label="处理摘要",
                    lines=15,
                    max_lines=30,
                    interactive=False,
                    elem_classes="result-box"
                )
        
        # 示例
        gr.Markdown("""
        ---
        ### 📖 使用说明
        
        | 步骤 | 操作 |
        |------|------|
        | 1 | 选择 AI 提供商并填入 API Key |
        | 2 | 上传一张或多张图片 |
        | 3 | (可选) 自定义提示词模板 |
        | 4 | 点击「开始处理」 |
        | 5 | 下载 ZIP 压缩包 |
        
        **ZIP 内容格式：**
				1.jpg  ← 原图
				1.txt  ← 提示词
				2.png
				2.txt...
				
				        
        **Telegram Bot:** 设置环境变量 `TELEGRAM_BOT_TOKEN` 后自动启动 Bot 服务
        """)
        
        # 绑定事件
        submit_btn.click(
            fn=analyze_images,
            inputs=[
                file_input, provider, api_key, base_url,
                model, custom_prompt, max_concurrent
            ],
            outputs=[output_file, output_summary],
        )
    
    return app


# =============== Telegram Bot 后台线程 ===============

def start_telegram_bot_thread():
    """在后台线程启动 Telegram Bot"""
    if not config.TELEGRAM_BOT_TOKEN:
        logger.info("未配置 TELEGRAM_BOT_TOKEN，跳过 Bot 启动")
        return
    
    def run():
        from telegram_bot import create_bot_app
        
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        bot_app = create_bot_app()
        logger.info("🤖 Telegram Bot 后台启动")
        bot_app.run_polling(drop_pending_updates=True)
    
    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    logger.info("Telegram Bot 线程已启动")


# =============== 启动 ===============

if __name__ == "__main__":
    # 启动 Telegram Bot（后台线程）
    start_telegram_bot_thread()
    
    # 启动 Gradio Web UI
    app = create_gradio_app()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
    )