from enum import Enum
import uuid
from sqlmodel import SQLModel, Field
from sqlalchemy import (
    Column,
    DateTime,
    func,
    JSON,
    UniqueConstraint,
    Enum as SQLAlchemyEnum,
)
from datetime import datetime
from typing import Optional


class SampleStatus(str, Enum):
    AVAILABLE = "AVAILABLE"
    ARCHIVED = "ARCHIVED"


class RegressionSample(SQLModel, table=True):
    __tablename__ = "regression_dataset"
    __table_args__ = (UniqueConstraint("job_id", name="uq_regression_dataset_job_id"),)

    id: Optional[int] = Field(default=None, primary_key=True)
    job_id: uuid.UUID = Field(foreign_key="validation_jobs.id", index=True)
    message: Optional[str] = Field(default=None, nullable=True)
    request: Optional[dict] = Field(default=None, nullable=True, sa_type=JSON)
    steps: Optional[list] = Field(default=None, nullable=True, sa_type=JSON)
    query_runs: Optional[list] = Field(default=None, nullable=True, sa_type=JSON)
    gold_answer: Optional[dict] = Field(default=None, nullable=True, sa_type=JSON)
    meta: Optional[dict] = Field(default=None, nullable=True, sa_type=JSON)
    created_at: Optional[datetime] = Field(
        default=None, sa_column=Column(DateTime, default=func.now())
    )
    status: SampleStatus = Field(
        default=SampleStatus.AVAILABLE, sa_column=SQLAlchemyEnum(SampleStatus)
    )