from sqlalchemy import BIGINT, FLOAT, TIMESTAMP, Enum, Index, Integer, JSON, String, Text, text
from sqlalchemy.orm import Mapped, mapped_column

from app.database import Base


class TrainTask(Base):
    __tablename__ = 'train_tasks'
    __table_args__ = (
        Index('idx_task_status', 'status'),
        {'mysql_comment': '模型训练任务表'},
    )

    id: Mapped[int] = mapped_column(BIGINT, primary_key=True, autoincrement=True)
    model_name: Mapped[str] = mapped_column(String(255), nullable=False, comment='模型名称')
    package_id: Mapped[int] = mapped_column(BIGINT, nullable=False, comment='训练集数据包ID')
    package_name: Mapped[str] = mapped_column(String(255), nullable=False, comment='训练集数据包名称')
    val_package_id: Mapped[int | None] = mapped_column(BIGINT, nullable=True, comment='验证集数据包ID')
    val_package_name: Mapped[str | None] = mapped_column(String(255), nullable=True, comment='验证集数据包名称')
    epoch_logs: Mapped[list | None] = mapped_column(JSON, nullable=True, comment='每轮训练日志')
    params: Mapped[dict] = mapped_column(JSON, nullable=False, comment='LSTM超参数')
    status: Mapped[str] = mapped_column(
        Enum('pending', 'running', 'completed', 'failed', 'cancelled', name='train_status_enum'),
        nullable=False,
        server_default=text("'pending'"),
        comment='训练状态',
    )
    progress: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text('0'), comment='进度 0-100')
    train_loss: Mapped[float | None] = mapped_column(FLOAT, nullable=True, comment='训练损失')
    val_loss: Mapped[float | None] = mapped_column(FLOAT, nullable=True, comment='验证损失')
    test_loss: Mapped[float | None] = mapped_column(FLOAT, nullable=True, comment='测试损失')
    error_msg: Mapped[str | None] = mapped_column(Text, nullable=True, comment='错误信息')
    is_saved: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text('0'), comment='是否已保存为模型')
    created_at: Mapped[str] = mapped_column(TIMESTAMP, nullable=False, server_default=text('CURRENT_TIMESTAMP'))
    updated_at: Mapped[str] = mapped_column(
        TIMESTAMP,
        nullable=False,
        server_default=text('CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP'),
    )


class SavedModel(Base):
    __tablename__ = 'saved_models'
    __table_args__ = (
        Index('idx_saved_task', 'task_id'),
        {'mysql_comment': '已保存模型表'},
    )

    id: Mapped[int] = mapped_column(BIGINT, primary_key=True, autoincrement=True)
    task_id: Mapped[int] = mapped_column(BIGINT, nullable=False, comment='来源训练任务ID')
    model_name: Mapped[str] = mapped_column(String(255), nullable=False, comment='模型名称')
    package_id: Mapped[int] = mapped_column(BIGINT, nullable=False, comment='数据包ID')
    package_name: Mapped[str] = mapped_column(String(255), nullable=False, comment='数据包名称')
    params: Mapped[dict] = mapped_column(JSON, nullable=False, comment='LSTM超参数')
    file_path: Mapped[str] = mapped_column(String(500), nullable=False, comment='模型文件路径')
    description: Mapped[str | None] = mapped_column(Text, nullable=True, comment='模型说明')
    train_loss: Mapped[float | None] = mapped_column(FLOAT, nullable=True)
    val_loss: Mapped[float | None] = mapped_column(FLOAT, nullable=True)
    test_loss: Mapped[float | None] = mapped_column(FLOAT, nullable=True)
    created_at: Mapped[str] = mapped_column(TIMESTAMP, nullable=False, server_default=text('CURRENT_TIMESTAMP'))
