import sqlalchemy as sa
from sqlalchemy.sql import func

from .enums import TaskStatus

_constaints_naming_convention = {
    "ix": "ix_%(column_0_label)s",
    "uq": "uq_%(table_name)s_%(column_0_name)s",
    "ck": "ck_%(table_name)s_%(constraint_name)s",
    "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
    "pk": "pk_%(table_name)s",
}

metadata = sa.MetaData(naming_convention=_constaints_naming_convention)


class NonNullable(sa.Column):
    def __init__(self, *args, **kwargs):
        kwargs["nullable"] = False
        super().__init__(*args, **kwargs)


tasks = sa.Table(
    "tasks",
    metadata,
    sa.Column("id", sa.BigInteger, primary_key=True),
    NonNullable("created", sa.DateTime(timezone=True), server_default=func.now()),
    sa.Column(
        "current_log_id", sa.BigInteger, sa.ForeignKey("tasks_log.id"), nullable=True
    ),
    sa.Column("status", sa.Enum(TaskStatus)),
    NonNullable("timing_from", sa.DateTime(timezone=True), unique=True),
    NonNullable("timing_to", sa.DateTime(timezone=True), unique=True),
    sa.CheckConstraint("timing_from < timing_to", name="timig_from_lt_timing_to"),
    sa.Index(
        "tasks_status_in_progress_btree_index",
        "status",
        postgresql_using="btree",
        postgresql_where=sa.column("status").in_(TaskStatus.in_progress(values=True)),
    ),
    sa.Index(
        "tasks_status_failed_btree_index",
        "status",
        postgresql_using="btree",
        postgresql_where=sa.column("status").is_(None),
    ),
)

tasks_log = sa.Table(
    "tasks_log",
    metadata,
    sa.Column("id", sa.BigInteger, primary_key=True),
    NonNullable("task_id", sa.BigInteger, sa.ForeignKey("tasks.id")),
    NonNullable("created", sa.DateTime(timezone=True), server_default=func.now()),
    NonNullable("executor_id", sa.String(length=128)),
    NonNullable("status", sa.Enum(TaskStatus)),
    sa.Column("execution_state", sa.JSON, nullable=True),
    sa.UniqueConstraint("executor_id", "status"),
    sa.Index(
        "tasks_log_status_normalizer_in_progress_btree_index",
        "status",
        postgresql_using="btree",
        postgresql_where=sa.column("status") == TaskStatus.accepted_by_normalizer.value,
    ),
    sa.Index(
        "tasks_log_status_charger_in_progress_btree_index",
        "status",
        postgresql_using="btree",
        postgresql_where=sa.column("status").in_(
            TaskStatus.charger_in_progress(values=True)
        ),
    ),
    sa.Index(
        "tasks_log_status_collector_in_progress_btree_index",
        "status",
        postgresql_using="btree",
        postgresql_where=sa.column("status") == TaskStatus.accepted_by_collector.value,
    ),
)

locks = sa.Table(
    "locks",
    metadata,
    sa.Column("id", sa.Integer, primary_key=True),
    NonNullable("name", sa.String(length=128)),
    sa.UniqueConstraint("name"),
)
