SupervisorAI/schemas/face_feature.py

"""
人脸特征值的Pydantic模型
用于数据验证和序列化
"""

from datetime import datetime
from typing import Optional, List, Dict, Any
from pydantic import BaseModel, Field, field_validator, ConfigDict, model_validator
from enum import IntEnum


# 枚举定义（与数据库模型一致）
class FeatureStatus(IntEnum):
    NOT_STARTED = 0
    PROCESSING = 1
    SUCCESS = 2
    FAILED = 3


# 基础模型
class FaceFeatureBase(BaseModel):
    """基础模型，包含所有字段"""
    person_id: int = Field(..., description="人员ID", gt=0)
    feature_type: Optional[int] = Field(None, description="模型版本", ge=0)
    feature_data: Optional[bytes] = Field(None, description="特征值（二进制）")
    pic_id: Optional[str] = Field(None, description="图片ID", max_length=255)
    status: FeatureStatus = Field(
        default=FeatureStatus.NOT_STARTED,
        description="计算状态"
    )
    start_time: Optional[datetime] = Field(None, description="计算开始时间")
    finish_time: Optional[datetime] = Field(None, description="计算结束时间")

    @field_validator('feature_data', mode='before')
    @classmethod
    def validate_feature_data(cls, v):
        """验证特征数据"""
        if v is not None and not isinstance(v, bytes):
            if isinstance(v, str):
                # 尝试从hex字符串转换
                try:
                    return bytes.fromhex(v)
                except ValueError:
                    raise ValueError("feature_data must be valid hex string or bytes")
            else:
                raise ValueError("feature_data must be bytes or hex string")
        return v


# 创建模型
class FaceFeatureCreate(FaceFeatureBase):
    """创建特征记录模型"""
    # 创建时不指定ID和时间
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "person_id": 1001,
                "feature_type": 1,
                "pic_id": "img_20250101_001",
                "status": 0
            }
        }
    )


class FaceFeatureUpdate(BaseModel):
    """更新特征记录模型"""
    feature_type: Optional[int] = Field(None, description="模型版本", ge=0)
    feature_data: Optional[bytes] = Field(None, description="特征值（二进制）")
    pic_id: Optional[str] = Field(None, description="图片ID", max_length=255)
    status: Optional[FeatureStatus] = Field(None, description="计算状态")
    start_time: Optional[datetime] = Field(None, description="计算开始时间")
    finish_time: Optional[datetime] = Field(None, description="计算结束时间")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "status": 2,
                "finish_time": "2024-01-01T12:00:00Z"
            }
        }
    )


# 查询参数模型
class FaceFeatureQuery(BaseModel):
    """特征记录查询参数"""
    person_id: Optional[int] = Field(None, description="人员ID", gt=0)
    feature_type: Optional[int] = Field(None, description="模型版本", ge=0)
    status: Optional[FeatureStatus] = Field(None, description="计算状态")
    start_date: Optional[datetime] = Field(None, description="开始时间")
    end_date: Optional[datetime] = Field(None, description="结束时间")
    has_feature_data: Optional[bool] = Field(None, description="是否有特征数据")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "person_id": 1001,
                "status": 2,
                "start_date": "2024-01-01T00:00:00Z"
            }
        }
    )


# 响应模型
class FaceFeatureResponse(FaceFeatureBase):
    """特征记录响应模型"""
    id: int
    created_time: datetime

    # 计算字段（将在验证后设置）
    status_name: Optional[str] = None
    is_completed: Optional[bool] = None
    processing_time: Optional[float] = None
    has_feature_data: Optional[bool] = None

    @model_validator(mode='after')
    def set_computed_fields(self):
        """设置所有计算字段"""
        # 状态名称
        try:
            self.status_name = FeatureStatus(self.status).name
        except ValueError:
            self.status_name = f"未知状态({self.status})"

        # 是否完成
        self.is_completed = self.status in [FeatureStatus.SUCCESS, FeatureStatus.FAILED]

        # 处理时间
        if self.start_time and self.finish_time:
            self.processing_time = (self.finish_time - self.start_time).total_seconds()

        # 是否有特征数据
        self.has_feature_data = self.feature_data is not None and len(self.feature_data) > 0

        return self

    model_config = ConfigDict(
        from_attributes=True,
        populate_by_name=True,
        json_schema_extra={
            "example": {
                "id": 1,
                "person_id": 1001,
                "feature_type": 1,
                "status": 2,
                "status_name": "SUCCESS",
                "created_time": "2024-01-01T10:00:00Z",
                "pic_id": "img_20250101_001",
                "start_time": "2024-01-01T10:00:00Z",
                "finish_time": "2024-01-01T10:00:05Z",
                "is_completed": True,
                "processing_time": 5.0,
                "has_feature_data": True
            }
        }
    )


class FaceFeatureListResponse(BaseModel):
    """特征记录列表响应"""
    total: int = Field(..., description="总记录数")
    items: List[FaceFeatureResponse] = Field(..., description="记录列表")
    page: Optional[int] = Field(None, description="当前页码")
    page_size: Optional[int] = Field(None, description="每页数量")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "total": 100,
                "page": 1,
                "page_size": 20,
                "items": []
            }
        }
    )


# 批量操作模型
class BatchFaceFeatureCreate(BaseModel):
    """批量创建特征记录"""
    items: List[FaceFeatureCreate] = Field(..., description="特征记录列表", min_items=1, max_items=1000)

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "items": [
                    {"person_id": 1001, "feature_type": 1},
                    {"person_id": 1002, "feature_type": 1}
                ]
            }
        }
    )


class FaceFeatureStatsResponse(BaseModel):
    """特征记录统计响应"""
    total_count: int = Field(..., description="总记录数")
    by_status: Dict[str, int] = Field(..., description="按状态统计")
    by_feature_type: Dict[str, int] = Field(..., description="按特征类型统计")
    avg_processing_time: Optional[float] = Field(None, description="平均处理时间（秒）")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "total_count": 1000,
                "by_status": {"SUCCESS": 800, "PROCESSING": 100, "FAILED": 100},
                "by_feature_type": {"1": 500, "2": 500},
                "avg_processing_time": 5.2
            }
        }
    )