processor.py

import os
import asyncio
import shutil
import zipfile
import logging
import time
from pathlib import Path
from typing import List, Optional, Callable, Tuple
from dataclasses import dataclass

from api_clients import get_ai_client, BaseClient
from config import config

logger = logging.getLogger(__name__)


@dataclass
class ProcessResult:
    index: int
    filename: str
    prompt: str
    success: bool
    error: Optional[str] = None


class BatchProcessor:
    """批量图片处理器"""
    
    def __init__(
        self,
        ai_client: Optional[BaseClient] = None,
        max_concurrent: int = None,
        custom_prompt: Optional[str] = None,
    ):
        self.client = ai_client or get_ai_client()
        self.max_concurrent = max_concurrent or config.MAX_CONCURRENT
        self.semaphore = asyncio.Semaphore(self.max_concurrent)
        self.custom_prompt = custom_prompt
    
    async def process_single(
        self,
        image_path: str,
        index: int,
        progress_callback: Optional[Callable] = None,
    ) -> ProcessResult:
        """处理单张图片"""
        filename = Path(image_path).name
        
        async with self.semaphore:
            try:
                logger.info(f"[{index}] 开始处理: {filename}")
                prompt = await self.client.analyze_image(
                    image_path, self.custom_prompt
                )
                result = ProcessResult(
                    index=index,
                    filename=filename,
                    prompt=prompt,
                    success=True,
                )
                logger.info(f"[{index}] 完成: {filename}")
                
                if progress_callback:
                    await progress_callback(index, filename, True, prompt)
                    
                return result
                
            except Exception as e:
                error_msg = str(e)
                logger.error(f"[{index}] 失败: {filename} - {error_msg}")
                result = ProcessResult(
                    index=index,
                    filename=filename,
                    prompt="",
                    success=False,
                    error=error_msg,
                )
                if progress_callback:
                    await progress_callback(index, filename, False, error_msg)
                return result
    
    async def process_batch(
        self,
        image_paths: List[str],
        output_dir: Optional[str] = None,
        progress_callback: Optional[Callable] = None,
    ) -> Tuple[str, List[ProcessResult]]:
        """
        批量处理图片
        返回: (zip文件路径, 处理结果列表)
        """
        # 创建输出目录
        timestamp = int(time.time())
        output_dir = output_dir or os.path.join(config.OUTPUT_DIR, f"batch_{timestamp}")
        os.makedirs(output_dir, exist_ok=True)
        
        # 并发处理
        tasks = []
        for i, img_path in enumerate(image_paths, 1):
            task = self.process_single(img_path, i, progress_callback)
            tasks.append(task)
        
        results = await asyncio.gather(*tasks)
        results.sort(key=lambda r: r.index)
        
        # 保存结果文件
        for result in results:
            if not result.success:
                continue
            
            src_path = image_paths[result.index - 1]
            suffix = Path(src_path).suffix or ".jpg"
            
            # 复制原图: 1.jpg, 2.png, ...
            dst_image = os.path.join(output_dir, f"{result.index}{suffix}")
            shutil.copy2(src_path, dst_image)
            
            # 保存提示词: 1.txt, 2.txt, ...
            dst_txt = os.path.join(output_dir, f"{result.index}.txt")
            with open(dst_txt, "w", encoding="utf-8") as f:
                f.write(result.prompt)
        
        # 打包 ZIP
        zip_path = f"{output_dir}.zip"
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for root, dirs, files in os.walk(output_dir):
                for file in sorted(files):
                    file_path = os.path.join(root, file)
                    arcname = file  # ZIP 内不含目录层级
                    zf.write(file_path, arcname)
        
        # 清理临时目录
        shutil.rmtree(output_dir, ignore_errors=True)
        
        success_count = sum(1 for r in results if r.success)
        fail_count = sum(1 for r in results if not r.success)
        logger.info(f"批量处理完成: 成功 {success_count}, 失败 {fail_count}")
        
        return zip_path, results


def process_batch_sync(
    image_paths: List[str],
    provider: str = None,
    api_key: str = None,
    base_url: str = None,
    model: str = None,
    custom_prompt: str = None,
    max_concurrent: int = None,
) -> Tuple[str, str]:
    """同步包装器（给 Gradio 用）"""
    client = get_ai_client(provider, api_key, base_url, model)
    processor = BatchProcessor(
        ai_client=client,
        max_concurrent=max_concurrent or config.MAX_CONCURRENT,
        custom_prompt=custom_prompt if custom_prompt else None,
    )
    
    loop = asyncio.new_event_loop()
    try:
        zip_path, results = loop.run_until_complete(
            processor.process_batch(image_paths)
        )
    finally:
        loop.close()
    
    # 生成摘要
    summary_lines = []
    for r in results:
        status = "✅" if r.success else "❌"
        preview = r.prompt[:100] + "..." if r.success and len(r.prompt) > 100 else r.prompt
        if not r.success:
            preview = f"Error: {r.error}"
        summary_lines.append(f"{status} [{r.index}] {r.filename}\n   {preview}")
    
    summary = "\n\n".join(summary_lines)
    return zip_path, summary