ocr_platform/table_line_generator/backend/api/batch.py

"""
批量处理 API
"""

from fastapi import APIRouter, BackgroundTasks, HTTPException
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any
from pathlib import Path
from loguru import logger

from services.batch_service import BatchProcessor
from models.schemas import TableStructure, ImageSize


router = APIRouter(prefix="/api/batch", tags=["batch"])


class FilePair(BaseModel):
    """文件对"""
    json_path: str = Field(description="JSON 文件路径")
    image_path: str = Field(description="图片文件路径")


class BatchProcessRequest(BaseModel):
    """批量处理请求"""
    template_name: str = Field(description="要应用的模板名称")
    file_pairs: List[FilePair] = Field(description="文件对列表")
    output_dir: str = Field(description="输出目录")
    parallel: bool = Field(default=True, description="是否并行处理")
    adjust_rows: bool = Field(default=True, description="是否自适应调整行分割")
    structure_suffix: str = Field(default="_structure.json", description="结构文件后缀")
    image_suffix: str = Field(default="_with_lines.png", description="输出图片文件后缀")


class BatchProcessResult(BaseModel):
    """单个文件处理结果"""
    success: bool
    json_path: str
    image_path: str
    structure_path: Optional[str] = None
    filename: str
    rows: Optional[int] = None
    cols: Optional[int] = None
    error: Optional[str] = None


class BatchProcessResponse(BaseModel):
    """批量处理响应"""
    success: bool
    total: int
    processed: int
    failed: int
    results: List[Dict[str, Any]]
    message: Optional[str] = None


class DrawBatchRequest(BaseModel):
    """批量绘图请求"""
    results: List[Dict[str, Any]] = Field(description="处理结果列表")
    line_width: int = Field(default=2, description="线条宽度")
    line_color: List[int] = Field(default=[0, 0, 0], description="线条颜色 RGB")


class DrawBatchResponse(BaseModel):
    """批量绘图响应"""
    success: bool
    total: int
    drawn: int
    results: List[Dict[str, Any]]
    message: Optional[str] = None


@router.post("/process", response_model=BatchProcessResponse)
async def batch_process(request: BatchProcessRequest):
    """
    批量处理文件
    
    将指定模板的结构应用到多个文件：
    - 复用列边界（竖线）
    - 自适应调整行分割（横线）
    - 支持并行处理
    - 自动处理不同图片尺寸的坐标映射
    """
    try:
        # 验证输出目录
        output_path = Path(request.output_dir)
        if not output_path.parent.exists():
            raise HTTPException(
                status_code=400, 
                detail=f"输出目录的父目录不存在: {output_path.parent}"
            )
        
        # 转换文件对格式
        file_pairs = [
            {
                'json_path': pair.json_path,
                'image_path': pair.image_path
            }
            for pair in request.file_pairs
        ]
        
        # 创建处理器
        processor = BatchProcessor(max_workers=4)
        
        # 执行批量处理
        summary = processor.process_batch_from_data_source(
            template_name=request.template_name,
            file_pairs=file_pairs,
            output_dir=request.output_dir,
            parallel=request.parallel,
            adjust_rows=request.adjust_rows,
            structure_suffix=request.structure_suffix,
            image_suffix=request.image_suffix
        )
        
        return BatchProcessResponse(
            success=True,
            total=summary['total'],
            processed=summary['success'],
            failed=summary['failed'],
            results=summary['results'],
            message=f"批量处理完成: 成功 {summary['success']}/{summary['total']}"
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.exception(f"批量处理失败: {e}")
        raise HTTPException(status_code=500, detail=f"批量处理失败: {e}")


@router.post("/draw", response_model=DrawBatchResponse)
async def batch_draw(request: DrawBatchRequest):
    """
    批量绘制表格线到图片上
    
    根据处理结果中的结构文件，在原图上绘制表格线
    """
    try:
        processor = BatchProcessor()
        
        # 确保颜色是 RGB 三元组
        line_color = tuple(request.line_color[:3]) if len(request.line_color) >= 3 else (0, 0, 0)
        
        draw_results = processor.draw_batch_images(
            results=request.results,
            line_width=request.line_width,
            line_color=line_color  # type: ignore
        )
        
        success_count = sum(1 for r in draw_results if r.get('success'))
        
        return DrawBatchResponse(
            success=True,
            total=len(request.results),
            drawn=success_count,
            results=draw_results,
            message=f"绘制完成: {success_count}/{len(request.results)} 张图片"
        )
        
    except Exception as e:
        logger.exception(f"批量绘图失败: {e}")
        raise HTTPException(status_code=500, detail=f"批量绘图失败: {e}")


@router.get("/health")
async def batch_health_check():
    """批量处理服务健康检查"""
    return {
        "status": "ok",
        "service": "batch-processor"
    }