import uuid
from datetime import datetime
from enum import Enum
from typing import Optional, Any

from pydantic import BaseModel
from sqlalchemy import Column, DateTime, func, JSON, Enum as SQLAlchemyEnum
from sqlmodel import Field, SQLModel


class Step(BaseModel):
    explanation: str
    sql: str


class ValidationJobStatus(str, Enum):
    CREATED = "CREATED"
    EXTRACTING = "EXTRACTING"
    VALIDATING = "VALIDATING"
    SUCCESS = "SUCCESS"
    FAILED = "FAILED"
    SKIPPED = "SKIPPED"
    EXPIRED = "EXPIRED"
    REPLAYED = "REPLAYED"


# class QueryRunStatus(Enum):
#     CREATED = "CREATED"
#     RUNNING = "RUNNING"
#     SUCCESS = "SUCCESS"
#     FAILED = "FAILED"


class QueryRun(BaseModel):
    explanation: Optional[str] = None
    sql: str
    # status: QueryRunStatus = QueryRunStatus.CREATED
    success: Optional[bool] = None
    message: Optional[str] = None
    results: Optional[list[list[Any]]] = []
    sql_error_code: Optional[int] = None
    sql_error_message: Optional[str] = None
    warnings: Optional[list[str]] = []
    llm_verification: Optional[str] = None


class ValidationJob(SQLModel, table=True):
    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        nullable=False,
    )

    status: ValidationJobStatus = Field(
        sa_column=Column(SQLAlchemyEnum(ValidationJobStatus))
    )
    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)
    instance_id: Optional[int] = Field(default=None, foreign_key="db_instances.id")
    external_request_id: Optional[str] = Field(
        max_length=128, default=None, nullable=True, unique=True
    )
    created_at: Optional[datetime] = Field(
        default=None, sa_column=Column(DateTime, default=func.now())
    )
    started_at: Optional[datetime] = Field(
        default=None, sa_column=Column(DateTime, default=func.now())
    )
    __tablename__ = "validation_jobs"