SupervisorAI/models/face_feature.py

"""
人脸特征数据模型
对应数据库表：sur_face_feature
"""

from typing import Optional
from datetime import datetime
from enum import IntEnum
from sqlalchemy import (
    Column,
    Integer,
    SmallInteger,
    LargeBinary,
    DateTime,
    Text,
    Index,
    UniqueConstraint
)
from sqlalchemy.dialects.postgresql import BYTEA
from sqlalchemy.sql import func

from database.base import BaseModel


class FeatureStatus(IntEnum):
    """人脸特征值计算状态枚举"""
    NOT_STARTED = 0  # 未开始
    PROCESSING = 1  # 计算中
    SUCCESS = 2  # 计算成功
    FAILED = 3  # 计算失败


# 导出别名以保持向后兼容性
FeatureStatusEnum = FeatureStatus


class SurFaceFeature(BaseModel):
    """
    人脸特征值表模型

    对应表结构：
    - id: 主键
    - person_id: 人员ID
    - feature_type: 模型版本
    - feature_data: 特征值（二进制）
    - created_time: 创建时间
    - pic_id: 图片ID
    - status: 计算状态
    - start_time: 计算开始时间
    - finish_time: 计算结束时间
    """

    __tablename__ = "sur_face_feature"
    __table_args__ = (
        # 唯一约束：person_id + feature_type
        UniqueConstraint("person_id", "feature_type", name="sur_face_feature_unique"),
        # 为常用查询字段创建索引
        Index("ix_sur_face_feature_person_id", "person_id"),
        Index("ix_sur_face_feature_feature_type", "feature_type"),
        Index("ix_sur_face_feature_status", "status"),
        Index("ix_sur_face_feature_created_time", "created_time"),
        {"schema": "public", "comment": "人脸特征值表"}
    )

    # 主键（自增序列）
    id = Column(
        Integer,
        primary_key=True,
        index=True,
        comment="主键"
    )

    # 人员ID（必填）
    person_id = Column(
        Integer,
        nullable=False,
        comment="人员id"
    )

    # 模型版本（特征类型）
    feature_type = Column(
        SmallInteger,
        nullable=True,  # 根据SQL，允许NULL
        comment="模型版本"
    )

    # 特征值（二进制数据）
    feature_data = Column(
        BYTEA,  # PostgreSQL的二进制类型
        nullable=True,
        comment="特征值"
    )

    # 创建时间（自动设置）
    created_time = Column(
        DateTime(timezone=True),
        server_default=func.now(),
        nullable=False,
        comment="创建时间"
    )

    # 图片ID
    pic_id = Column(
        Text,
        nullable=True,
        comment="图片id"
    )

    # 计算状态
    status = Column(
        SmallInteger,
        default=FeatureStatusEnum.NOT_STARTED,
        nullable=False,
        comment="人脸特征值计算状态：0=未开始，1=计算中，2=计算成功，3=计算失败"
    )

    # 计算开始时间
    start_time = Column(
        DateTime(timezone=True),
        nullable=True,
        comment="特征计算开始时间"
    )

    # 计算结束时间
    finish_time = Column(
        DateTime(timezone=True),
        nullable=True,
        comment="特征计算结束时间"
    )

    # 属性方法
    @property
    def status_name(self) -> str:
        """获取状态名称"""
        try:
            return FeatureStatusEnum(self.status).name
        except ValueError:
            return f"未知状态({self.status})"

    @property
    def is_completed(self) -> bool:
        """是否完成计算"""
        return self.status in [FeatureStatusEnum.SUCCESS, FeatureStatusEnum.FAILED]

    @property
    def processing_time(self) -> Optional[float]:
        """计算处理时间（秒）"""
        if self.start_time and self.finish_time:
            return (self.finish_time - self.start_time).total_seconds()
        return None

    def start_processing(self) -> None:
        """开始处理"""
        self.status = FeatureStatusEnum.PROCESSING
        self.start_time = datetime.now()
        self.finish_time = None

    def finish_processing(self, success: bool = True) -> None:
        """结束处理"""
        self.status = FeatureStatusEnum.SUCCESS if success else FeatureStatusEnum.FAILED
        self.finish_time = datetime.now()

    def to_dict(self, exclude: list = None) -> dict:
        """
        重写to_dict方法，处理二进制数据

        Args:
            exclude: 要排除的字段列表

        Returns:
            转换后的字典
        """
        exclude = exclude or []

        # 默认排除二进制数据（太大）
        default_exclude = ["feature_data"]
        final_exclude = list(set(exclude + default_exclude))

        result = super().to_dict(final_exclude)

        # 添加额外属性
        result["status_name"] = self.status_name
        result["is_completed"] = self.is_completed
        result["processing_time"] = self.processing_time

        # 如果有feature_data，添加一个标识
        if self.feature_data and "feature_data" not in exclude:
            result["has_feature_data"] = True
            result["feature_data_length"] = len(self.feature_data)
        else:
            result["has_feature_data"] = False

        return result

    def __repr__(self) -> str:
        return (f"<SurFaceFeature(id={self.id}, person_id={self.person_id}, "
                f"feature_type={self.feature_type}, status={self.status_name})>")