"""add_fair_share_and_resource_usage_history_tables

Revision ID: 71343531dd5a
Revises: 7369d1eb7d4a
Create Date: 2026-01-14 02:39:55.482831

"""

import sqlalchemy as sa
from alembic import op

from ai.backend.manager.models.base import GUID, ResourceSlotColumn

# revision identifiers, used by Alembic.
revision = "71343531dd5a"
down_revision = "7369d1eb7d4a"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "domain_fair_shares",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("scaling_group", sa.String(length=64), nullable=False),
        sa.Column("domain_name", sa.String(length=64), nullable=False),
        sa.Column(
            "weight",
            sa.Numeric(precision=10, scale=4),
            nullable=False,
            comment="Priority weight multiplier. Higher weight = higher priority allocation ratio. Example: 2.0 means 2x priority compared to default 1.0",
        ),
        sa.Column(
            "fair_share_factor",
            sa.Numeric(precision=8, scale=6),
            nullable=False,
            comment="Calculated priority score from 0.0 to 1.0. Higher value = less past usage = higher scheduling priority. Formula: F = 2^(-normalized_usage / weight)",
        ),
        sa.Column(
            "total_decayed_usage",
            ResourceSlotColumn(),
            nullable=False,
            comment="Sum of historical resource usage with time decay applied. Unit: resource-seconds (e.g., cpu-seconds, mem-seconds). Older usage is weighted less via half-life decay.",
        ),
        sa.Column(
            "resource_weights",
            ResourceSlotColumn(),
            nullable=False,
            comment="Resource weights used in fair share calculation. From scheduler_opts. Example: {cpu: 1.0, mem: 1.0, cuda.device: 10.0}",
        ),
        sa.Column(
            "normalized_usage",
            sa.Numeric(precision=8, scale=6),
            nullable=False,
            comment="Weighted average of (usage/capacity) per resource (0.0 ~ 1.0). Both usage and capacity are in resource-seconds. Formula: sum((usage[r]/capacity[r]) * weight[r]) / sum(weight[r])",
        ),
        sa.Column(
            "last_calculated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
            comment="Timestamp when fair_share_factor was last recalculated by batch job.",
        ),
        sa.Column(
            "lookback_start",
            sa.Date(),
            server_default=sa.text("CURRENT_DATE"),
            nullable=False,
            comment="Start date of the usage history window used in last calculation. Typically: today - lookback_days.",
        ),
        sa.Column(
            "lookback_end",
            sa.Date(),
            server_default=sa.text("CURRENT_DATE"),
            nullable=False,
            comment="End date of the usage history window used in last calculation. Typically: today.",
        ),
        sa.Column(
            "half_life_days",
            sa.Integer(),
            nullable=False,
            comment="Number of days for usage to decay to 50%%. Example: 7 means usage from 7 days ago counts as 50%% of recent usage.",
        ),
        sa.Column(
            "lookback_days",
            sa.Integer(),
            nullable=False,
            comment="Number of days of usage history to consider for fair share calculation. Example: 28 means only last 28 days of usage affects priority.",
        ),
        sa.Column(
            "decay_unit_days",
            sa.Integer(),
            nullable=False,
            comment="Aggregation period for usage buckets in days. Example: 1 means daily aggregation, 7 means weekly aggregation.",
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_domain_fair_shares")),
        sa.UniqueConstraint("scaling_group", "domain_name", name="uq_domain_fair_share"),
    )
    op.create_index(
        "ix_domain_fair_share_lookup",
        "domain_fair_shares",
        ["scaling_group", "domain_name"],
        unique=False,
    )
    op.create_index(
        op.f("ix_domain_fair_shares_domain_name"),
        "domain_fair_shares",
        ["domain_name"],
        unique=False,
    )
    op.create_index(
        op.f("ix_domain_fair_shares_scaling_group"),
        "domain_fair_shares",
        ["scaling_group"],
        unique=False,
    )
    op.create_table(
        "domain_usage_buckets",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("domain_name", sa.String(length=64), nullable=False),
        sa.Column("scaling_group", sa.String(length=64), nullable=False),
        sa.Column("period_start", sa.Date(), nullable=False),
        sa.Column("period_end", sa.Date(), nullable=False),
        sa.Column("decay_unit_days", sa.Integer(), nullable=False),
        sa.Column("resource_usage", ResourceSlotColumn(), nullable=False),
        sa.Column(
            "capacity_snapshot",
            ResourceSlotColumn(),
            nullable=False,
            comment="Scaling group capacity at bucket period. Sum of agent.available_slots for calculating usage ratio.",
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_domain_usage_buckets")),
        sa.UniqueConstraint(
            "domain_name", "scaling_group", "period_start", name="uq_domain_usage_bucket"
        ),
    )
    op.create_index(
        "ix_domain_usage_bucket_lookup",
        "domain_usage_buckets",
        ["domain_name", "scaling_group", "period_start"],
        unique=False,
    )
    op.create_index(
        op.f("ix_domain_usage_buckets_domain_name"),
        "domain_usage_buckets",
        ["domain_name"],
        unique=False,
    )
    op.create_table(
        "kernel_usage_records",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("kernel_id", GUID(), nullable=False),
        sa.Column("session_id", GUID(), nullable=False),
        sa.Column("user_uuid", GUID(), nullable=False),
        sa.Column("project_id", GUID(), nullable=False),
        sa.Column("domain_name", sa.String(length=64), nullable=False),
        sa.Column("scaling_group", sa.String(length=64), nullable=False),
        sa.Column("period_start", sa.DateTime(timezone=True), nullable=False),
        sa.Column("period_end", sa.DateTime(timezone=True), nullable=False),
        sa.Column("resource_usage", ResourceSlotColumn(), nullable=False),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_kernel_usage_records")),
    )
    op.create_index(
        op.f("ix_kernel_usage_records_domain_name"),
        "kernel_usage_records",
        ["domain_name"],
        unique=False,
    )
    op.create_index(
        op.f("ix_kernel_usage_records_kernel_id"),
        "kernel_usage_records",
        ["kernel_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_kernel_usage_records_project_id"),
        "kernel_usage_records",
        ["project_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_kernel_usage_records_scaling_group"),
        "kernel_usage_records",
        ["scaling_group"],
        unique=False,
    )
    op.create_index(
        op.f("ix_kernel_usage_records_user_uuid"),
        "kernel_usage_records",
        ["user_uuid"],
        unique=False,
    )
    op.create_index(
        "ix_kernel_usage_sg_period",
        "kernel_usage_records",
        ["scaling_group", "period_start"],
        unique=False,
    )
    op.create_index(
        "ix_kernel_usage_user_period",
        "kernel_usage_records",
        ["user_uuid", "period_start"],
        unique=False,
    )
    op.create_table(
        "project_fair_shares",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("scaling_group", sa.String(length=64), nullable=False),
        sa.Column("project_id", GUID(), nullable=False),
        sa.Column("domain_name", sa.String(length=64), nullable=False),
        sa.Column(
            "weight",
            sa.Numeric(precision=10, scale=4),
            nullable=False,
            comment="Priority weight multiplier. Higher weight = higher priority allocation ratio. Example: 2.0 means 2x priority compared to default 1.0",
        ),
        sa.Column(
            "fair_share_factor",
            sa.Numeric(precision=8, scale=6),
            nullable=False,
            comment="Calculated priority score from 0.0 to 1.0. Higher value = less past usage = higher scheduling priority. Formula: F = 2^(-normalized_usage / weight)",
        ),
        sa.Column(
            "total_decayed_usage",
            ResourceSlotColumn(),
            nullable=False,
            comment="Sum of historical resource usage with time decay applied. Unit: resource-seconds.",
        ),
        sa.Column(
            "resource_weights",
            ResourceSlotColumn(),
            nullable=False,
            comment="Resource weights used in fair share calculation. From scheduler_opts. Example: {cpu: 1.0, mem: 1.0, cuda.device: 10.0}",
        ),
        sa.Column(
            "normalized_usage",
            sa.Numeric(precision=8, scale=6),
            nullable=False,
            comment="Weighted average of (usage/capacity) per resource (0.0 ~ 1.0). Both usage and capacity are in resource-seconds. Formula: sum((usage[r]/capacity[r]) * weight[r]) / sum(weight[r])",
        ),
        sa.Column(
            "last_calculated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
            comment="Timestamp when fair_share_factor was last recalculated.",
        ),
        sa.Column(
            "lookback_start",
            sa.Date(),
            server_default=sa.text("CURRENT_DATE"),
            nullable=False,
            comment="Start date of the usage history window used in last calculation.",
        ),
        sa.Column(
            "lookback_end",
            sa.Date(),
            server_default=sa.text("CURRENT_DATE"),
            nullable=False,
            comment="End date of the usage history window used in last calculation.",
        ),
        sa.Column(
            "half_life_days",
            sa.Integer(),
            nullable=False,
            comment="Number of days for usage to decay to 50%%.",
        ),
        sa.Column(
            "lookback_days",
            sa.Integer(),
            nullable=False,
            comment="Number of days of usage history to consider.",
        ),
        sa.Column(
            "decay_unit_days",
            sa.Integer(),
            nullable=False,
            comment="Aggregation period for usage buckets in days.",
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_project_fair_shares")),
        sa.UniqueConstraint("scaling_group", "project_id", name="uq_project_fair_share"),
    )
    op.create_index(
        "ix_project_fair_share_lookup",
        "project_fair_shares",
        ["scaling_group", "project_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_project_fair_shares_domain_name"),
        "project_fair_shares",
        ["domain_name"],
        unique=False,
    )
    op.create_index(
        op.f("ix_project_fair_shares_project_id"),
        "project_fair_shares",
        ["project_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_project_fair_shares_scaling_group"),
        "project_fair_shares",
        ["scaling_group"],
        unique=False,
    )
    op.create_table(
        "project_usage_buckets",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("project_id", GUID(), nullable=False),
        sa.Column("domain_name", sa.String(length=64), nullable=False),
        sa.Column("scaling_group", sa.String(length=64), nullable=False),
        sa.Column("period_start", sa.Date(), nullable=False),
        sa.Column("period_end", sa.Date(), nullable=False),
        sa.Column("decay_unit_days", sa.Integer(), nullable=False),
        sa.Column("resource_usage", ResourceSlotColumn(), nullable=False),
        sa.Column(
            "capacity_snapshot",
            ResourceSlotColumn(),
            nullable=False,
            comment="Scaling group capacity at bucket period. Sum of agent.available_slots for calculating usage ratio.",
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_project_usage_buckets")),
        sa.UniqueConstraint(
            "project_id", "scaling_group", "period_start", name="uq_project_usage_bucket"
        ),
    )
    op.create_index(
        "ix_project_usage_bucket_lookup",
        "project_usage_buckets",
        ["project_id", "scaling_group", "period_start"],
        unique=False,
    )
    op.create_index(
        op.f("ix_project_usage_buckets_domain_name"),
        "project_usage_buckets",
        ["domain_name"],
        unique=False,
    )
    op.create_index(
        op.f("ix_project_usage_buckets_project_id"),
        "project_usage_buckets",
        ["project_id"],
        unique=False,
    )
    op.create_table(
        "user_fair_shares",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("scaling_group", sa.String(length=64), nullable=False),
        sa.Column("user_uuid", GUID(), nullable=False),
        sa.Column("project_id", GUID(), nullable=False),
        sa.Column("domain_name", sa.String(length=64), nullable=False),
        sa.Column(
            "weight",
            sa.Numeric(precision=10, scale=4),
            nullable=False,
            comment="Priority weight multiplier. Higher weight = higher priority allocation ratio. Example: 2.0 means 2x priority compared to default 1.0",
        ),
        sa.Column(
            "fair_share_factor",
            sa.Numeric(precision=8, scale=6),
            nullable=False,
            comment="Calculated priority score from 0.0 to 1.0. Higher value = less past usage = higher scheduling priority. Formula: F = 2^(-normalized_usage / weight)",
        ),
        sa.Column(
            "total_decayed_usage",
            ResourceSlotColumn(),
            nullable=False,
            comment="Sum of historical resource usage with time decay applied. Unit: resource-seconds.",
        ),
        sa.Column(
            "resource_weights",
            ResourceSlotColumn(),
            nullable=False,
            comment="Resource weights used in fair share calculation. From scheduler_opts. Example: {cpu: 1.0, mem: 1.0, cuda.device: 10.0}",
        ),
        sa.Column(
            "normalized_usage",
            sa.Numeric(precision=8, scale=6),
            nullable=False,
            comment="Weighted average of (usage/capacity) per resource (0.0 ~ 1.0). Both usage and capacity are in resource-seconds. Formula: sum((usage[r]/capacity[r]) * weight[r]) / sum(weight[r])",
        ),
        sa.Column(
            "last_calculated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
            comment="Timestamp when fair_share_factor was last recalculated.",
        ),
        sa.Column(
            "lookback_start",
            sa.Date(),
            server_default=sa.text("CURRENT_DATE"),
            nullable=False,
            comment="Start date of the usage history window used in last calculation.",
        ),
        sa.Column(
            "lookback_end",
            sa.Date(),
            server_default=sa.text("CURRENT_DATE"),
            nullable=False,
            comment="End date of the usage history window used in last calculation.",
        ),
        sa.Column(
            "half_life_days",
            sa.Integer(),
            nullable=False,
            comment="Number of days for usage to decay to 50%%.",
        ),
        sa.Column(
            "lookback_days",
            sa.Integer(),
            nullable=False,
            comment="Number of days of usage history to consider.",
        ),
        sa.Column(
            "decay_unit_days",
            sa.Integer(),
            nullable=False,
            comment="Aggregation period for usage buckets in days.",
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_user_fair_shares")),
        sa.UniqueConstraint("scaling_group", "user_uuid", "project_id", name="uq_user_fair_share"),
    )
    op.create_index(
        "ix_user_fair_share_lookup",
        "user_fair_shares",
        ["scaling_group", "user_uuid", "project_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_user_fair_shares_domain_name"), "user_fair_shares", ["domain_name"], unique=False
    )
    op.create_index(
        op.f("ix_user_fair_shares_project_id"), "user_fair_shares", ["project_id"], unique=False
    )
    op.create_index(
        op.f("ix_user_fair_shares_scaling_group"),
        "user_fair_shares",
        ["scaling_group"],
        unique=False,
    )
    op.create_index(
        op.f("ix_user_fair_shares_user_uuid"), "user_fair_shares", ["user_uuid"], unique=False
    )
    op.create_table(
        "user_usage_buckets",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("user_uuid", GUID(), nullable=False),
        sa.Column("project_id", GUID(), nullable=False),
        sa.Column("domain_name", sa.String(length=64), nullable=False),
        sa.Column("scaling_group", sa.String(length=64), nullable=False),
        sa.Column("period_start", sa.Date(), nullable=False),
        sa.Column("period_end", sa.Date(), nullable=False),
        sa.Column("decay_unit_days", sa.Integer(), nullable=False),
        sa.Column("resource_usage", ResourceSlotColumn(), nullable=False),
        sa.Column(
            "capacity_snapshot",
            ResourceSlotColumn(),
            nullable=False,
            comment="Scaling group capacity at bucket period. Sum of agent.available_slots for calculating usage ratio.",
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_user_usage_buckets")),
        sa.UniqueConstraint(
            "user_uuid", "project_id", "scaling_group", "period_start", name="uq_user_usage_bucket"
        ),
    )
    op.create_index(
        "ix_user_usage_bucket_lookup",
        "user_usage_buckets",
        ["user_uuid", "project_id", "scaling_group", "period_start"],
        unique=False,
    )
    op.create_index(
        op.f("ix_user_usage_buckets_domain_name"),
        "user_usage_buckets",
        ["domain_name"],
        unique=False,
    )
    op.create_index(
        op.f("ix_user_usage_buckets_project_id"), "user_usage_buckets", ["project_id"], unique=False
    )
    op.create_index(
        op.f("ix_user_usage_buckets_user_uuid"), "user_usage_buckets", ["user_uuid"], unique=False
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f("ix_user_usage_buckets_user_uuid"), table_name="user_usage_buckets")
    op.drop_index(op.f("ix_user_usage_buckets_project_id"), table_name="user_usage_buckets")
    op.drop_index(op.f("ix_user_usage_buckets_domain_name"), table_name="user_usage_buckets")
    op.drop_index("ix_user_usage_bucket_lookup", table_name="user_usage_buckets")
    op.drop_table("user_usage_buckets")
    op.drop_index(op.f("ix_user_fair_shares_user_uuid"), table_name="user_fair_shares")
    op.drop_index(op.f("ix_user_fair_shares_scaling_group"), table_name="user_fair_shares")
    op.drop_index(op.f("ix_user_fair_shares_project_id"), table_name="user_fair_shares")
    op.drop_index(op.f("ix_user_fair_shares_domain_name"), table_name="user_fair_shares")
    op.drop_index("ix_user_fair_share_lookup", table_name="user_fair_shares")
    op.drop_table("user_fair_shares")
    op.drop_index(op.f("ix_project_usage_buckets_project_id"), table_name="project_usage_buckets")
    op.drop_index(op.f("ix_project_usage_buckets_domain_name"), table_name="project_usage_buckets")
    op.drop_index("ix_project_usage_bucket_lookup", table_name="project_usage_buckets")
    op.drop_table("project_usage_buckets")
    op.drop_index(op.f("ix_project_fair_shares_scaling_group"), table_name="project_fair_shares")
    op.drop_index(op.f("ix_project_fair_shares_project_id"), table_name="project_fair_shares")
    op.drop_index(op.f("ix_project_fair_shares_domain_name"), table_name="project_fair_shares")
    op.drop_index("ix_project_fair_share_lookup", table_name="project_fair_shares")
    op.drop_table("project_fair_shares")
    op.drop_index("ix_kernel_usage_user_period", table_name="kernel_usage_records")
    op.drop_index("ix_kernel_usage_sg_period", table_name="kernel_usage_records")
    op.drop_index(op.f("ix_kernel_usage_records_user_uuid"), table_name="kernel_usage_records")
    op.drop_index(op.f("ix_kernel_usage_records_scaling_group"), table_name="kernel_usage_records")
    op.drop_index(op.f("ix_kernel_usage_records_project_id"), table_name="kernel_usage_records")
    op.drop_index(op.f("ix_kernel_usage_records_kernel_id"), table_name="kernel_usage_records")
    op.drop_index(op.f("ix_kernel_usage_records_domain_name"), table_name="kernel_usage_records")
    op.drop_table("kernel_usage_records")
    op.drop_index(op.f("ix_domain_usage_buckets_domain_name"), table_name="domain_usage_buckets")
    op.drop_index("ix_domain_usage_bucket_lookup", table_name="domain_usage_buckets")
    op.drop_table("domain_usage_buckets")
    op.drop_index(op.f("ix_domain_fair_shares_scaling_group"), table_name="domain_fair_shares")
    op.drop_index(op.f("ix_domain_fair_shares_domain_name"), table_name="domain_fair_shares")
    op.drop_index("ix_domain_fair_share_lookup", table_name="domain_fair_shares")
    op.drop_table("domain_fair_shares")
    # ### end Alembic commands ###
