from datetime import datetime as DateTime
from typing import TypeAlias

import sqlalchemy as sa
from databricks.sqlalchemy import TIMESTAMP
from sqlalchemy.orm import (
    DeclarativeBase,
    declared_attr,
    Mapped,
    mapped_column,
    MappedAsDataclass,
)
from sqlalchemy.types import (
    Double,
    String,
)
from typing_extensions import Annotated


float64: TypeAlias = Double
uint64: TypeAlias = sa.BigInteger
str32 = Annotated[str, 32]
str64 = Annotated[str, 64]

TABLE_NAME_MAPPING = {
    'Model': 'model_metadata',
}


class Base(MappedAsDataclass, DeclarativeBase):

    @declared_attr.directive
    def __tablename__(cls) -> str:
        return TABLE_NAME_MAPPING[cls.__name__]

    type_annotation_map = {
        float64: Double,
        uint64: sa.BigInteger,
        str32: String(32),
        str64: String(64),
    }

    insert_timestamp: Mapped[DateTime] = mapped_column(
        TIMESTAMP,
        init=False,
        nullable=False,
        server_default=sa.func.current_timestamp(),
    )


class Model(Base):
    __table_args__ = (
        sa.UniqueConstraint("name", "version", name='unique_key'),
    )
    pkid: Mapped[uint64] = mapped_column(
        sa.BigInteger,
        sa.Identity(always=True),
        init=False,
        primary_key=True,
    )
    name: Mapped[str64]
    version: Mapped[str32]