ocr_platform/table_line_generator/backend/models/schemas.py

"""
Pydantic 数据模型定义
"""

from pydantic import BaseModel, Field
from typing import List, Dict, Optional, Any


class TableStructure(BaseModel):
    """表格结构"""
    horizontal_lines: List[int] = Field(default_factory=list, description="横线Y坐标列表")
    vertical_lines: List[int] = Field(default_factory=list, description="竖线X坐标列表")
    table_bbox: Optional[List[int]] = Field(default=None, description="表格边界框 [x1, y1, x2, y2]")
    row_height: Optional[int] = Field(default=None, description="标准行高")
    col_widths: Optional[List[int]] = Field(default=None, description="列宽列表")
    total_rows: Optional[int] = Field(default=None, description="总行数")
    total_cols: Optional[int] = Field(default=None, description="总列数")
    mode: Optional[str] = Field(default="hybrid", description="分析模式")
    modified_h_lines: List[int] = Field(default_factory=list, description="被修改的横线索引")
    modified_v_lines: List[int] = Field(default_factory=list, description="被修改的竖线索引")
    image_rotation_angle: float = Field(default=0.0, description="图片旋转角度")
    skew_angle: float = Field(default=0.0, description="倾斜角度")
    is_skew_corrected: bool = Field(default=False, description="是否已校正倾斜")


class ImageSize(BaseModel):
    """图片尺寸"""
    width: int
    height: int


class UploadResponse(BaseModel):
    """上传响应"""
    success: bool
    image_base64: str = Field(description="Base64 编码的图片")
    structure: TableStructure = Field(description="表格结构")
    image_size: ImageSize = Field(description="图片尺寸")
    scale_factor: float = Field(default=1.0, description="缩放比例 (1.0 表示未缩放)")
    ocr_data: Optional[Dict[str, Any]] = Field(default=None, description="原始 OCR 数据 (用于重新分析)")
    suggested_filename: Optional[str] = Field(default=None, description="建议的文件名")
    message: Optional[str] = None


class AnalyzeParams(BaseModel):
    """分析参数"""
    y_tolerance: int = Field(default=5, ge=1, le=50, description="Y轴聚类容差")
    x_tolerance: int = Field(default=10, ge=1, le=100, description="X轴聚类容差")
    min_row_height: int = Field(default=20, ge=5, le=100, description="最小行高")
    method: str = Field(default="auto", description="分析方法: auto/cluster/mineru")


class AnalyzeRequest(BaseModel):
    """重新分析请求"""
    ocr_data: Dict[str, Any] = Field(description="OCR 数据")
    params: AnalyzeParams = Field(default_factory=AnalyzeParams)


class AnalyzeResponse(BaseModel):
    """分析响应"""
    success: bool
    structure: TableStructure
    message: Optional[str] = None


class SaveRequest(BaseModel):
    """保存请求"""
    structure: TableStructure = Field(description="表格结构")
    image_base64: Optional[str] = Field(default=None, description="Base64 编码的原始图片（可选，saveImage=false时不需要）")
    output_dir: str = Field(description="输出目录路径")
    filename: str = Field(description="结构文件名 (不含扩展名)")
    image_filename: Optional[str] = Field(default=None, description="图片文件名 (不含扩展名)")
    overwrite_mode: str = Field(default="overwrite", description="覆盖策略: overwrite/skip/new")
    structure_suffix: str = Field(default="_structure", description="结构文件后缀（不含.json扩展名）")
    image_suffix: str = Field(default="", description="图片文件后缀（不含.png扩展名）")
    line_width: int = Field(default=2, description="线条宽度")
    line_color: List[int] = Field(default=[0, 0, 0], description="线条颜色 RGB")


class SaveResponse(BaseModel):
    """保存响应"""
    success: bool
    structure_path: Optional[str] = Field(default=None, description="结构文件路径")
    image_path: Optional[str] = Field(default=None, description="图片文件路径（可选）")
    message: Optional[str] = None


class ErrorResponse(BaseModel):
    """错误响应"""
    success: bool = False
    error: str
    detail: Optional[str] = None


class LoadByPathRequest(BaseModel):
    """按路径加载请求"""
    json_path: str = Field(description="JSON 文件路径")
    image_path: str = Field(description="图片文件路径")
    output_dir: Optional[str] = Field(default=None, description="输出目录路径（可选，用于查找已保存的标注结构）")
    structure_suffix: str = Field(default="_structure", description="结构文件后缀（不含.json）")
    image_suffix: str = Field(default="", description="图片文件后缀（不含.png）")


class PreviewApplyRequest(BaseModel):
    """预览应用模板请求"""
    template_name: str = Field(description="模板名称")
    target_image_size: ImageSize = Field(description="目标图片尺寸")
    target_table_bbox: Optional[List[int]] = Field(default=None, description="目标表格边界框 [x1, y1, x2, y2]")


class ApplyTemplateRequest(BaseModel):
    """应用模板请求"""
    template_name: str = Field(description="模板名称")
    structure: TableStructure = Field(description="原始表格结构")
    target_table_bbox: Optional[List[int]] = Field(default=None, description="目标表格边界框")


class CreateTemplateRequest(BaseModel):
    """创建模板请求"""
    name: str = Field(description="模板名称")
    description: Optional[str] = Field(default=None, description="模板描述")
    structure: TableStructure = Field(description="表格结构")


class TemplateInfo(BaseModel):
    """模板信息"""
    name: str
    description: Optional[str]
    created_at: str
    row_count: int
    col_count: int


class TemplateListResponse(BaseModel):
    """模板列表响应"""
    success: bool
    templates: List[TemplateInfo]
    message: Optional[str] = None


class TemplateResponse(BaseModel):
    """单个模板响应"""
    success: bool
    data: Dict
    message: Optional[str] = None


class DeleteTemplateRequest(BaseModel):
    """删除模板请求"""
    name: str = Field(description="模板名称")


class DeleteTemplateResponse(BaseModel):
    """删除模板响应"""
    success: bool
    message: str


class CreateTemplateResponse(BaseModel):
    """创建模板响应"""
    success: bool
    data: Dict = Field(description="包含 name, path, stats, created_at")
    message: Optional[str] = None


class GetTemplateResponse(BaseModel):
    """获取模板响应"""
    success: bool
    data: Dict = Field(description="模板完整信息")
    message: Optional[str] = None


class TemplateApplyResponse(BaseModel):
    """预览/应用模板响应"""
    success: bool
    data: Dict = Field(description="应用后的表格结构")
    message: Optional[str] = None


class HealthResponse(BaseModel):
    """健康检查响应"""
    status: str = Field(description="服务状态")
    service: str = Field(description="服务名称")


class HomeDirectoryResponse(BaseModel):
    """主目录响应"""
    path: str = Field(description="用户主目录路径")