#
# Copyright (c) 2021 the Hugging Face team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.#
"""create_main_tables
Revision ID: 2dc2179b353f
Revises: 
Create Date: 2021-04-28 07:54:49.181680
"""  # noqa
from typing import Tuple

import sqlalchemy as sa
from alembic import op


# revision identifiers, used by Alembic
revision = "2dc2179b353f"
down_revision = None
branch_labels = None
depends_on = None


def create_updated_at_trigger() -> None:
    op.execute(
        """
        CREATE OR REPLACE FUNCTION update_updated_at_column()
            RETURNS TRIGGER AS
        $$
        BEGIN
            NEW.updated_at = now();
            RETURN NEW;
        END;
        $$ language 'plpgsql';
        """
    )


def timestamps(indexed: bool = False) -> Tuple[sa.Column, sa.Column]:
    return (
        sa.Column(
            "created_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.func.now(),
            nullable=False,
            index=indexed,
        ),
        sa.Column(
            "updated_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.func.now(),
            nullable=False,
            index=indexed,
        ),
    )


def create_experiments_table() -> None:
    op.create_table(
        "experiments",
        sa.Column("id", sa.Integer, primary_key=True),
        sa.Column("name", sa.Text, nullable=False, index=True),
        sa.Column("owner", sa.Text, nullable=False, index=True),
        *timestamps(),
    )
    op.execute(
        """
        CREATE TRIGGER update_experiments_modtime
            BEFORE UPDATE
            ON experiments
            FOR EACH ROW
        EXECUTE PROCEDURE update_updated_at_column();
        """
    )


def create_whitelist_table() -> None:
    op.create_table(
        "whitelist",
        sa.Column("id", sa.Integer, primary_key=True),
        sa.Column("experiment_id", sa.Integer, nullable=False, index=True),
        sa.Column("user_id", sa.Integer, nullable=False, index=False),
        *timestamps(),
    )
    op.execute(
        """
        CREATE TRIGGER update_whitelist_modtime
            BEFORE UPDATE
            ON whitelist
            FOR EACH ROW
        EXECUTE PROCEDURE update_updated_at_column();
        """
    )


def create_users_table() -> None:
    op.create_table(
        "users",
        sa.Column("id", sa.Integer, primary_key=True),
        sa.Column("username", sa.Text, unique=True, nullable=True, index=True),
        *timestamps(),
    )
    op.execute(
        """
        CREATE TRIGGER update_user_modtime
            BEFORE UPDATE
            ON users
            FOR EACH ROW
        EXECUTE PROCEDURE update_updated_at_column();
        """
    )


def upgrade() -> None:
    create_updated_at_trigger()
    create_experiments_table()
    create_whitelist_table()
    create_users_table()


def downgrade() -> None:
    op.drop_table("users")
    op.drop_table("experiments")
    op.drop_table("whitelist")
    op.execute("DROP FUNCTION update_updated_at_column")
