from typing import Type

from ktem.db.engine import engine
from sqlalchemy import JSON, Boolean, Column, String
from sqlalchemy.orm import DeclarativeBase
from theflow.settings import settings as flowsettings
from theflow.utils.modules import import_dotted_string


class Base(DeclarativeBase):
    pass


class BaseRerankingTable(Base):
    """Base table to store rerankings model"""

    __abstract__ = True

    name = Column(String, primary_key=True, unique=True)
    spec = Column(JSON, default={})
    default = Column(Boolean, default=False)


__base_reranking: Type[BaseRerankingTable] = (
    import_dotted_string(flowsettings.KH_TABLE_RERANKING, safe=False)
    if hasattr(flowsettings, "KH_TABLE_RERANKING")
    else BaseRerankingTable
)


class RerankingTable(__base_reranking):  # type: ignore
    __tablename__ = "reranking"


if not getattr(flowsettings, "KH_ENABLE_ALEMBIC", False):
    RerankingTable.metadata.create_all(engine)
