telegram_bot.py

import os
import asyncio
import logging
import tempfile
import time
from pathlib import Path
from typing import Dict, List
from dataclasses import dataclass, field

from telegram import Update, InputFile, BotCommand
from telegram.ext import (
    Application, CommandHandler, MessageHandler, 
    filters, ContextTypes, ConversationHandler
)

from api_clients import get_ai_client
from processor import BatchProcessor, ProcessResult
from config import config

logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    level=logging.INFO
)
logger = logging.getLogger(__name__)


@dataclass
class UserSession:
    """用户会话，用于收集批量图片"""
    images: List[str] = field(default_factory=list)
    collecting: bool = False
    start_time: float = 0
    custom_prompt: str = ""
    provider: str = ""


# 用户会话存储
user_sessions: Dict[int, UserSession] = {}


def get_session(user_id: int) -> UserSession:
    if user_id not in user_sessions:
        user_sessions[user_id] = UserSession()
    return user_sessions[user_id]


def check_auth(user_id: int) -> bool:
    """检查用户权限"""
    if not config.ALLOWED_USERS:
        return True  # 未配置白名单则允许所有人
    return user_id in config.ALLOWED_USERS


# =============== 命令处理 ===============

async def cmd_start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """开始命令"""
    if not check_auth(update.effective_user.id):
        await update.message.reply_text("⛔ 无权限使用此 Bot")
        return
    
    welcome = """🖼️ **图片反推提示词 Bot**

📌 **使用方法：**

**单张处理：** 直接发送一张图片

**批量处理：**
1. /batch - 开始收集图片
2. 连续发送多张图片
3. /done - 完成收集并开始处理

**其他命令：**
/prompt <text> - 自定义提示词模板
/provider <name> - 切换 AI (openai/gemini/claude/qwen)
/status - 查看当前设置
/cancel - 取消当前批次
/help - 帮助信息

发送图片即可开始！"""
    
    await update.message.reply_text(welcome, parse_mode="Markdown")


async def cmd_help(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if not check_auth(update.effective_user.id):
        return
    
    help_text = """📖 **帮助**

**支持格式：** JPG, PNG, WEBP, GIF, BMP

**输出格式：**
• 单张：直接回复提示词文本
• 批量：输出 ZIP 压缩包
  - 1.jpg + 1.txt
  - 2.png + 2.txt
  - ...

**AI 提供商：**
• `openai` - GPT-4o (默认)
• `gemini` - Google Gemini
• `claude` - Anthropic Claude
• `qwen` - 通义千问VL
• `custom` - 自定义 OpenAI 兼容 API

**自定义提示词：**
`/prompt Describe this image in detail for Stable Diffusion`

**重置提示词：**
`/prompt reset`"""
    
    await update.message.reply_text(help_text, parse_mode="Markdown")


async def cmd_batch(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """开始批量收集"""
    if not check_auth(update.effective_user.id):
        return
    
    user_id = update.effective_user.id
    session = get_session(user_id)
    
    # 清理旧图片
    for img in session.images:
        if os.path.exists(img):
            os.remove(img)
    
    session.images = []
    session.collecting = True
    session.start_time = time.time()
    
    await update.message.reply_text(
        "📥 **批量模式已开启**\n\n"
        "请连续发送图片，完成后发送 /done 开始处理\n"
        f"最多支持 {config.MAX_BATCH_SIZE} 张\n"
        "发送 /cancel 取消",
        parse_mode="Markdown"
    )


async def cmd_done(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """完成收集，开始处理"""
    if not check_auth(update.effective_user.id):
        return
    
    user_id = update.effective_user.id
    session = get_session(user_id)
    
    if not session.collecting:
        await update.message.reply_text("❌ 当前没有在收集图片，请先 /batch")
        return
    
    if not session.images:
        await update.message.reply_text("❌ 还没有收到任何图片")
        return
    
    session.collecting = False
    count = len(session.images)
    
    status_msg = await update.message.reply_text(
        f"⏳ 开始处理 {count} 张图片...\n进度: 0/{count}"
    )
    
    # 进度回调
    completed = [0]
    
    async def progress_cb(index, filename, success, content):
        completed[0] += 1
        status = "✅" if success else "❌"
        try:
            await status_msg.edit_text(
                f"⏳ 处理中...\n"
                f"进度: {completed[0]}/{count}\n"
                f"最新: {status} {filename}"
            )
        except Exception:
            pass
    
    # 处理
    try:
        provider = session.provider or config.AI_PROVIDER
        client = get_ai_client(provider=provider)
        processor = BatchProcessor(
            ai_client=client,
            custom_prompt=session.custom_prompt or None,
        )
        
        zip_path, results = await processor.process_batch(
            session.images, progress_callback=progress_cb
        )
        
        success_count = sum(1 for r in results if r.success)
        fail_count = sum(1 for r in results if not r.success)
        
        # 发送结果摘要
        summary = f"✅ **处理完成**\n\n"
        summary += f"• 成功: {success_count}\n"
        summary += f"• 失败: {fail_count}\n\n"
        
        # 显示每个结果的简短预览
        for r in results:
            status = "✅" if r.success else "❌"
            if r.success:
                preview = r.prompt[:80].replace("\n", " ")
                summary += f"{status} `{r.index}`: {preview}...\n"
            else:
                summary += f"{status} `{r.index}`: {r.error}\n"
        
        await status_msg.edit_text(summary, parse_mode="Markdown")
        
        # 发送 ZIP
        if os.path.exists(zip_path):
            with open(zip_path, "rb") as f:
                await update.message.reply_document(
                    document=InputFile(f, filename=f"prompts_{int(time.time())}.zip"),
                    caption=f"📦 {success_count} 张图片的提示词结果"
                )
            os.remove(zip_path)
    
    except Exception as e:
        logger.error(f"批量处理失败: {e}")
        await status_msg.edit_text(f"❌ 处理失败: {str(e)}")
    
    finally:
        # 清理
        for img in session.images:
            if os.path.exists(img):
                os.remove(img)
        session.images = []


async def cmd_cancel(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """取消批量收集"""
    if not check_auth(update.effective_user.id):
        return
    
    user_id = update.effective_user.id
    session = get_session(user_id)
    
    for img in session.images:
        if os.path.exists(img):
            os.remove(img)
    
    session.images = []
    session.collecting = False
    
    await update.message.reply_text("❌ 已取消批量收集")


async def cmd_prompt(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """设置自定义提示词"""
    if not check_auth(update.effective_user.id):
        return
    
    user_id = update.effective_user.id
    session = get_session(user_id)
    
    text = " ".join(context.args) if context.args else ""
    
    if not text:
        current = session.custom_prompt or "默认提示词"
        await update.message.reply_text(
            f"📝 当前提示词:\n\n`{current}`\n\n"
            f"用法: `/prompt <你的提示词>` 或 `/prompt reset`",
            parse_mode="Markdown"
        )
        return
    
    if text.lower() == "reset":
        session.custom_prompt = ""
        await update.message.reply_text("✅ 已重置为默认提示词")
    else:
        session.custom_prompt = text
        await update.message.reply_text(f"✅ 提示词已更新:\n\n`{text}`", parse_mode="Markdown")


async def cmd_provider(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """切换 AI 提供商"""
    if not check_auth(update.effective_user.id):
        return
    
    user_id = update.effective_user.id
    session = get_session(user_id)
    
    valid = ["openai", "gemini", "claude", "qwen", "custom"]
    
    if not context.args:
        current = session.provider or config.AI_PROVIDER
        await update.message.reply_text(
            f"🤖 当前 AI: `{current}`\n\n"
            f"可选: {', '.join(valid)}\n"
            f"用法: `/provider openai`",
            parse_mode="Markdown"
        )
        return
    
    provider = context.args[0].lower()
    if provider not in valid:
        await update.message.reply_text(f"❌ 不支持: {provider}\n可选: {', '.join(valid)}")
        return
    
    session.provider = provider
    await update.message.reply_text(f"✅ 已切换为: `{provider}`", parse_mode="Markdown")


async def cmd_status(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """查看当前设置"""
    if not check_auth(update.effective_user.id):
        return
    
    user_id = update.effective_user.id
    session = get_session(user_id)
    
    provider = session.provider or config.AI_PROVIDER
    prompt_preview = (session.custom_prompt[:80] + "...") if session.custom_prompt else "默认"
    batch_status = f"收集中 ({len(session.images)} 张)" if session.collecting else "未开启"
    
    await update.message.reply_text(
        f"⚙️ **当前设置**\n\n"
        f"• AI 提供商: `{provider}`\n"
        f"• 提示词: {prompt_preview}\n"
        f"• 批量模式: {batch_status}\n"
        f"• 并发数: {config.MAX_CONCURRENT}\n"
        f"• 最大批量: {config.MAX_BATCH_SIZE}",
        parse_mode="Markdown"
    )


# =============== 图片处理 ===============

async def handle_photo(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """处理接收到的图片"""
    if not check_auth(update.effective_user.id):
        await update.message.reply_text("⛔ 无权限")
        return
    
    user_id = update.effective_user.id
    session = get_session(user_id)
    
    # 下载图片
    if update.message.document:
        file = await update.message.document.get_file()
        suffix = Path(update.message.document.file_name).suffix or ".jpg"
    elif update.message.photo:
        file = await update.message.photo[-1].get_file()  # 最高分辨率
        suffix = ".jpg"
    else:
        return
    
    # 保存到临时文件
    tmp_dir = tempfile.mkdtemp()
    tmp_path = os.path.join(tmp_dir, f"img_{int(time.time()*1000)}{suffix}")
    await file.download_to_drive(tmp_path)
    
    # 批量模式：收集图片
    if session.collecting:
        if len(session.images) >= config.MAX_BATCH_SIZE:
            await update.message.reply_text(
                f"⚠️ 已达上限 {config.MAX_BATCH_SIZE} 张，请发送 /done 开始处理"
            )
            return
        
        session.images.append(tmp_path)
        await update.message.reply_text(
            f"📥 已收到第 {len(session.images)} 张图片\n"
            f"继续发送或 /done 开始处理"
        )
        return
    
    # 单张模式：立即处理
    processing_msg = await update.message.reply_text("⏳ 正在分析图片...")
    
    try:
        provider = session.provider or config.AI_PROVIDER
        client = get_ai_client(provider=provider)
        prompt = await client.analyze_image(
            tmp_path,
            session.custom_prompt or None
        )
        
        # 回复提示词
        await processing_msg.edit_text(
            f"✅ **提示词结果：**\n\n```\n{prompt}\n```",
            parse_mode="Markdown"
        )
        
        # 同时发送 txt 文件
        txt_path = tmp_path.replace(suffix, ".txt")
        with open(txt_path, "w", encoding="utf-8") as f:
            f.write(prompt)
        
        with open(txt_path, "rb") as f:
            await update.message.reply_document(
                document=InputFile(f, filename="prompt.txt"),
                caption="📄 提示词文件"
            )
        
        os.remove(txt_path)
    
    except Exception as e:
        logger.error(f"处理失败: {e}")
        await processing_msg.edit_text(f"❌ 处理失败: {str(e)}")
    
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        os.rmdir(os.path.dirname(tmp_path))


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

async def post_init(application):
    """设置 Bot 命令菜单"""
    commands = [
        BotCommand("start", "开始使用"),
        BotCommand("batch", "开始批量收集图片"),
        BotCommand("done", "完成收集并处理"),
        BotCommand("cancel", "取消批量收集"),
        BotCommand("prompt", "设置/查看自定义提示词"),
        BotCommand("provider", "切换AI提供商"),
        BotCommand("status", "查看当前设置"),
        BotCommand("help", "帮助信息"),
    ]
    await application.bot.set_my_commands(commands)


def create_bot_app() -> Application:
    """创建 Bot Application"""
    app = (
        Application.builder()
        .token(config.TELEGRAM_BOT_TOKEN)
        .post_init(post_init)
        .build()
    )
    
    # 注册处理器
    app.add_handler(CommandHandler("start", cmd_start))
    app.add_handler(CommandHandler("help", cmd_help))
    app.add_handler(CommandHandler("batch", cmd_batch))
    app.add_handler(CommandHandler("done", cmd_done))
    app.add_handler(CommandHandler("cancel", cmd_cancel))
    app.add_handler(CommandHandler("prompt", cmd_prompt))
    app.add_handler(CommandHandler("provider", cmd_provider))
    app.add_handler(CommandHandler("status", cmd_status))
    
    # 图片处理（压缩图 + 文件形式）
    app.add_handler(MessageHandler(
        filters.PHOTO | (filters.Document.IMAGE), 
        handle_photo
    ))
    
    return app


def run_bot():
    """独立运行 Telegram Bot"""
    if not config.TELEGRAM_BOT_TOKEN:
        logger.error("未设置 TELEGRAM_BOT_TOKEN")
        return
    
    logger.info("Telegram Bot 启动中...")
    app = create_bot_app()
    app.run_polling(drop_pending_updates=True)


if __name__ == "__main__":
    run_bot()