from pydantic import BaseModel
from typing import Dict, Optional
from datetime import datetime


class ExtractionStage(BaseModel):
    time: int
    status: str
    variation: str


class ExtractionRecordBase(BaseModel):
    id: int
    fileName: str
    fileType: str
    fileSize: str
    extractedAt: datetime
    status: str
    confidence: float
    fieldsExtracted: int
    totalTime: int
    stages: Dict[str, ExtractionStage]
    errorMessage: Optional[str] = None

    class Config:
        orm_mode = True