OCR/backend/app/schemas/task.py

"""
Tool_OCR - Task Management Schemas
"""

from typing import Optional, List
from datetime import datetime
from pydantic import BaseModel, Field
from enum import Enum


class TaskStatusEnum(str, Enum):
    """Task status enumeration"""
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


class ProcessingTrackEnum(str, Enum):
    """Processing track enumeration for dual-track processing"""
    OCR = "ocr"           # PaddleOCR PP-StructureV3 for scanned documents
    DIRECT = "direct"     # PyMuPDF for editable PDFs
    HYBRID = "hybrid"     # Mixed processing
    AUTO = "auto"         # Auto-detect best track


class LayoutModelEnum(str, Enum):
    """Layout detection model selection for OCR track.

    Different models are optimized for different document types:
    - CHINESE: PP-DocLayout-S, optimized for Chinese documents (forms, contracts, invoices)
    - DEFAULT: PubLayNet-based, optimized for English academic papers
    - CDLA: CDLA model, specialized Chinese document layout analysis
    """
    CHINESE = "chinese"   # PP-DocLayout-S - Best for Chinese documents (recommended)
    DEFAULT = "default"   # PubLayNet-based - Best for English documents
    CDLA = "cdla"         # CDLA model - Alternative for Chinese layout


class TaskCreate(BaseModel):
    """Task creation request"""
    filename: Optional[str] = Field(None, description="Original filename")
    file_type: Optional[str] = Field(None, description="File MIME type")


class TaskUpdate(BaseModel):
    """Task update request"""
    status: Optional[TaskStatusEnum] = None
    error_message: Optional[str] = None
    processing_time_ms: Optional[int] = None
    result_json_path: Optional[str] = None
    result_markdown_path: Optional[str] = None
    result_pdf_path: Optional[str] = None


class TaskFileResponse(BaseModel):
    """Task file response schema"""
    id: int
    original_name: Optional[str] = None
    stored_path: Optional[str] = None
    file_size: Optional[int] = None
    mime_type: Optional[str] = None
    file_hash: Optional[str] = None
    created_at: datetime

    class Config:
        from_attributes = True


class TaskResponse(BaseModel):
    """Task response schema"""
    id: int
    user_id: int
    task_id: str
    filename: Optional[str] = None
    file_type: Optional[str] = None
    status: TaskStatusEnum
    result_json_path: Optional[str] = None
    result_markdown_path: Optional[str] = None
    result_pdf_path: Optional[str] = None
    error_message: Optional[str] = None
    processing_time_ms: Optional[int] = None
    created_at: datetime
    updated_at: datetime
    completed_at: Optional[datetime] = None
    file_deleted: bool = False

    class Config:
        from_attributes = True


class TaskDetailResponse(TaskResponse):
    """Detailed task response with files"""
    files: List[TaskFileResponse] = []
    # Dual-track processing field (extracted from result metadata)
    processing_track: Optional[ProcessingTrackEnum] = None


class TaskListResponse(BaseModel):
    """Paginated task list response"""
    tasks: List[TaskResponse]
    total: int
    page: int
    page_size: int
    has_more: bool


class TaskStatsResponse(BaseModel):
    """User task statistics"""
    total: int
    pending: int
    processing: int
    completed: int
    failed: int


class TaskHistoryQuery(BaseModel):
    """Task history query parameters"""
    status: Optional[TaskStatusEnum] = None
    filename: Optional[str] = None
    date_from: Optional[datetime] = None
    date_to: Optional[datetime] = None
    page: int = Field(default=1, ge=1)
    page_size: int = Field(default=50, ge=1, le=100)
    order_by: str = Field(default="created_at")
    order_desc: bool = Field(default=True)


class UploadFileInfo(BaseModel):
    """Uploaded file information"""
    filename: str
    file_size: int
    file_type: str


class UploadResponse(BaseModel):
    """File upload response"""
    task_id: str = Field(..., description="Created task ID")
    filename: str = Field(..., description="Original filename")
    file_size: int = Field(..., description="File size in bytes")
    file_type: str = Field(..., description="File MIME type")
    status: TaskStatusEnum = Field(..., description="Initial task status")


# ===== Dual-Track Processing Schemas =====

class PPStructureV3Params(BaseModel):
    """PP-StructureV3 fine-tuning parameters for OCR track.

    DEPRECATED: This class is deprecated and will be removed in a future version.
    Use `layout_model` parameter in ProcessingOptions instead.
    """
    layout_detection_threshold: Optional[float] = Field(
        None, ge=0, le=1,
        description="Layout block detection score threshold (lower=more blocks, higher=high confidence only)"
    )
    layout_nms_threshold: Optional[float] = Field(
        None, ge=0, le=1,
        description="Layout NMS IoU threshold (lower=aggressive overlap removal, higher=allow more overlap)"
    )
    layout_merge_bboxes_mode: Optional[str] = Field(
        None, pattern="^(union|large|small)$",
        description="Bbox merging strategy: 'small'=conservative, 'large'=aggressive, 'union'=middle"
    )
    layout_unclip_ratio: Optional[float] = Field(
        None, gt=0,
        description="Layout bbox expansion ratio (larger=looser boxes, smaller=tighter boxes)"
    )
    text_det_thresh: Optional[float] = Field(
        None, ge=0, le=1,
        description="Text detection score threshold (lower=detect more small/low-contrast text, higher=cleaner)"
    )
    text_det_box_thresh: Optional[float] = Field(
        None, ge=0, le=1,
        description="Text box candidate threshold (lower=more text boxes, higher=fewer false positives)"
    )
    text_det_unclip_ratio: Optional[float] = Field(
        None, gt=0,
        description="Text box expansion ratio (larger=looser boxes, smaller=tighter boxes)"
    )


class ProcessingOptions(BaseModel):
    """Processing options for dual-track OCR"""
    use_dual_track: bool = Field(default=True, description="Enable dual-track processing")
    force_track: Optional[ProcessingTrackEnum] = Field(None, description="Force specific track (ocr/direct)")
    language: str = Field(default="ch", description="OCR language code")
    include_layout: bool = Field(default=True, description="Include layout analysis")
    include_images: bool = Field(default=True, description="Extract and save images")
    confidence_threshold: Optional[float] = Field(None, ge=0, le=1, description="OCR confidence threshold")

    # Layout model selection (OCR track only)
    layout_model: Optional[LayoutModelEnum] = Field(
        default=LayoutModelEnum.CHINESE,
        description="Layout detection model: 'chinese' (recommended for Chinese docs), 'default' (English docs), 'cdla' (Chinese layout)"
    )


class AnalyzeRequest(BaseModel):
    """Document analysis request"""
    use_dual_track: bool = Field(default=True, description="Enable dual-track processing")
    force_track: Optional[ProcessingTrackEnum] = Field(None, description="Force specific track")
    language: str = Field(default="ch", description="OCR language")
    include_layout: bool = Field(default=True, description="Include layout analysis")


class DocumentAnalysisResponse(BaseModel):
    """Document type analysis response"""
    task_id: str
    filename: str
    recommended_track: ProcessingTrackEnum
    confidence: float = Field(..., ge=0, le=1, description="Detection confidence")
    reason: str = Field(..., description="Reason for recommendation")
    document_info: dict = Field(default_factory=dict, description="Document metadata")
    is_editable: bool = Field(..., description="Whether document has extractable text")
    text_coverage: Optional[float] = Field(None, description="Percentage of text coverage")
    page_count: Optional[int] = Field(None, description="Number of pages")


class ProcessingMetadata(BaseModel):
    """Processing metadata included in responses"""
    processing_track: ProcessingTrackEnum
    processing_time_seconds: float
    language: str
    page_count: int
    total_elements: int
    total_text_regions: int
    total_tables: int
    total_images: int
    average_confidence: Optional[float] = None
    unified_format: bool = True


class TaskResponseWithMetadata(TaskResponse):
    """Extended task response with processing metadata"""
    processing_track: Optional[ProcessingTrackEnum] = None
    processing_metadata: Optional[ProcessingMetadata] = None


class ExportOptions(BaseModel):
    """Export format options"""
    format: str = Field(default="json", description="Export format: json, markdown, pdf, unified")
    include_metadata: bool = Field(default=True, description="Include processing metadata")
    include_statistics: bool = Field(default=True, description="Include document statistics")
    legacy_format: bool = Field(default=False, description="Use legacy JSON format for compatibility")