"""create job metric table

Revision ID: 99c3877d0f10
Revises: 815022877cfe
Create Date: 2024-11-28 16:20:20.533111

"""

import time
from datetime import datetime, timezone
from textwrap import dedent

import sqlalchemy as sa
from sqlalchemy.types import TypeDecorator, DateTime

from alembic import op

# revision identifiers, used by Alembic.
revision = "99c3877d0f10"
down_revision = "815022877cfe"
branch_labels = None
depends_on = None


create_materialized_view_template = dedent(
    """
    CREATE MATERIALIZED VIEW {view_name}
    WITH (timescaledb.continuous) AS
    SELECT job_submission_id,
            node_host,
            time_bucket('{time_bucket}', time) AS bucket,
            AVG(cpu_frequency) AS cpu_frequency,
            SUM(cpu_time) AS cpu_time,
            AVG(cpu_utilization) AS cpu_utilization,
            AVG(gpu_memory) AS gpu_memory,
            AVG(gpu_utilization) AS gpu_utilization,
            SUM(page_faults) AS page_faults,
            AVG(memory_rss) AS memory_rss,
            AVG(memory_virtual) AS memory_virtual,
            SUM(disk_read) AS disk_read,
            SUM(disk_write) AS disk_write
    FROM job_submission_metrics
    GROUP BY job_submission_id, node_host, bucket
    WITH NO DATA
    """
)

add_continuous_aggregate_policy_template = dedent(
    """
    SELECT add_continuous_aggregate_policy('{mat_view_name}',
        start_offset => INTERVAL '{start_offset}',
        end_offset => INTERVAL '{end_offset}',
        schedule_interval => INTERVAL '{schedule_interval}',
        initial_start => '{initial_start}')
    """
)

remove_continuous_aggregate_policy_template = "SELECT remove_continuous_aggregate_policy('{mat_view_name}')"

drop_materialized_view_template = "DROP MATERIALIZED VIEW {view_name}"


class TimestampInt(TypeDecorator):
    impl = DateTime(timezone=True)

    def process_bind_param(self, value: int | None, dialect) -> datetime | None:
        if value is not None:
            return datetime.fromtimestamp(value, tz=timezone.utc)
        return value

    def process_result_value(self, value: datetime | None, dialect) -> int | None:
        if value is not None:
            return int(value.timestamp())
        return value


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.execute(sa.text("CREATE EXTENSION IF NOT EXISTS timescaledb"))
    op.create_table(
        "job_submission_metrics",
        sa.Column(
            "time",
            TimestampInt(),
            nullable=False,
        ),
        sa.Column("job_submission_id", sa.Integer(), nullable=False),
        sa.Column("slurm_job_id", sa.Integer(), nullable=False),
        sa.Column("node_host", sa.String(), nullable=False),
        sa.Column("step", sa.Integer(), nullable=False),
        sa.Column("task", sa.Integer(), nullable=False),
        sa.Column("cpu_frequency", sa.Float(), nullable=False),
        sa.Column("cpu_time", sa.Float(), nullable=False),
        sa.Column("cpu_utilization", sa.Float(), nullable=False),
        sa.Column("gpu_memory", sa.BigInteger(), nullable=False),
        sa.Column("gpu_utilization", sa.Float(), nullable=False),
        sa.Column("page_faults", sa.BigInteger(), nullable=False),
        sa.Column("memory_rss", sa.BigInteger(), nullable=False),
        sa.Column("memory_virtual", sa.BigInteger(), nullable=False),
        sa.Column("disk_read", sa.BigInteger(), nullable=False),
        sa.Column("disk_write", sa.BigInteger(), nullable=False),
        sa.PrimaryKeyConstraint("time", "job_submission_id", "node_host", "step", "task"),
        sa.ForeignKeyConstraint(["job_submission_id"], ["job_submissions.id"], ondelete="CASCADE"),
    )
    op.create_index(
        "idx_job_submission_metrics_time",
        "job_submission_metrics",
        ["time"],
        unique=False,
    )
    op.create_index(
        "idx_node_host_step_task",
        "job_submission_metrics",
        ["node_host", "step", "task"],
        unique=False,
    )
    op.create_index(
        op.f("ix_job_submission_metrics_node_host"),
        "job_submission_metrics",
        ["node_host"],
        unique=False,
    )
    op.create_index(
        op.f("ix_job_submission_metrics_step"),
        "job_submission_metrics",
        ["step"],
        unique=False,
    )
    op.create_index(
        op.f("ix_job_submission_metrics_task"),
        "job_submission_metrics",
        ["task"],
        unique=False,
    )
    op.execute(sa.text("SELECT create_hypertable('job_submission_metrics', by_range('time'))"))
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_10_seconds_by_node",
                time_bucket="10 seconds",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_10_seconds_by_node",
                start_offset="20 seconds",
                end_offset="0 seconds",
                schedule_interval="10 seconds",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_10_seconds_all_nodes",
                time_bucket="10 seconds",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_10_seconds_all_nodes",
                start_offset="20 seconds",
                end_offset="0 seconds",
                schedule_interval="10 seconds",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_1_minute_by_node",
                time_bucket="1 minute",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_minute_by_node",
                start_offset="2 minutes",
                end_offset="0 seconds",
                schedule_interval="1 minute",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_1_minute_all_nodes",
                time_bucket="1 minute",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_minute_all_nodes",
                start_offset="2 minutes",
                end_offset="0 seconds",
                schedule_interval="1 minute",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_10_minutes_by_node",
                time_bucket="10 minutes",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_10_minutes_by_node",
                start_offset="20 minutes",
                end_offset="0 seconds",
                schedule_interval="10 minutes",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_10_minutes_all_nodes",
                time_bucket="10 minutes",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_10_minutes_all_nodes",
                start_offset="20 minutes",
                end_offset="0 seconds",
                schedule_interval="10 minutes",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_1_hour_by_node",
                time_bucket="1 hour",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_hour_by_node",
                start_offset="2 hours",
                end_offset="0 seconds",
                schedule_interval="1 hour",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_1_hour_all_nodes",
                time_bucket="1 hour",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_hour_all_nodes",
                start_offset="2 hours",
                end_offset="0 seconds",
                schedule_interval="1 hour",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_1_week_by_node",
                time_bucket="1 week",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_week_by_node",
                start_offset="2 weeks",
                end_offset="0 seconds",
                schedule_interval="1 week",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    op.execute(
        sa.text(
            create_materialized_view_template.format(
                view_name="metrics_nodes_mv_1_week_all_nodes",
                time_bucket="1 week",
            )
        )
    )
    op.execute(
        sa.text(
            add_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_week_all_nodes",
                start_offset="2 weeks",
                end_offset="0 seconds",
                schedule_interval="1 week",
                initial_start=datetime.now(timezone.utc).isoformat(),
            )
        )
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f("ix_job_submission_metrics_task"), table_name="job_submission_metrics")
    op.drop_index(op.f("ix_job_submission_metrics_step"), table_name="job_submission_metrics")
    op.drop_index(op.f("ix_job_submission_metrics_node_host"), table_name="job_submission_metrics")
    op.drop_index("idx_node_host_step_task", table_name="job_submission_metrics")
    op.drop_index("idx_job_submission_metrics_time", table_name="job_submission_metrics")
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_week_by_node"
            )
        )
    )
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_hour_by_node"
            )
        )
    )
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_10_minutes_by_node"
            )
        )
    )
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_minute_by_node"
            )
        )
    )
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_10_seconds_by_node"
            )
        )
    )
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_week_all_nodes"
            )
        )
    )
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_hour_all_nodes"
            )
        )
    )
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_10_minutes_all_nodes"
            )
        )
    )
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_1_minute_all_nodes"
            )
        )
    )
    op.execute(
        sa.text(
            remove_continuous_aggregate_policy_template.format(
                mat_view_name="metrics_nodes_mv_10_seconds_all_nodes"
            )
        )
    )
    op.execute(sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_1_week_by_node")))
    op.execute(sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_1_hour_by_node")))
    op.execute(
        sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_10_minutes_by_node"))
    )
    op.execute(sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_1_minute_by_node")))
    op.execute(
        sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_10_seconds_by_node"))
    )
    op.execute(sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_1_week_all_nodes")))
    op.execute(sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_1_hour_all_nodes")))
    op.execute(
        sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_10_minutes_all_nodes"))
    )
    op.execute(
        sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_1_minute_all_nodes"))
    )
    op.execute(
        sa.text(drop_materialized_view_template.format(view_name="metrics_nodes_mv_10_seconds_all_nodes"))
    )
    op.drop_table("job_submission_metrics")
    # ### end Alembic commands ###
