"""Repurpose JobUsageStatistics for new Steps table

Revision ID: 416b9d2db10b
Revises: 25fc99e97839
Create Date: 2025-01-17 11:27:42.115755

"""

from typing import Sequence, Union

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op
from letta.settings import settings

# revision identifiers, used by Alembic.
revision: str = "416b9d2db10b"
down_revision: Union[str, None] = "25fc99e97839"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Skip this migration for SQLite
    if not settings.letta_pg_uri_no_default:
        return

    # ### commands auto generated by Alembic - please adjust! ###
    # Rename the table
    op.rename_table("job_usage_statistics", "steps")

    # Rename the foreign key constraint and drop non-null constraint
    op.alter_column("steps", "job_id", nullable=True)
    op.drop_constraint("fk_job_usage_statistics_job_id", "steps", type_="foreignkey")

    # Change id field from int to string
    op.execute("ALTER TABLE steps RENAME COLUMN id TO old_id")
    op.add_column("steps", sa.Column("id", sa.String(), nullable=True))
    op.execute("""UPDATE steps SET id = 'step-' || gen_random_uuid()::text""")
    op.drop_column("steps", "old_id")
    op.alter_column("steps", "id", nullable=False)
    op.create_primary_key("pk_steps_id", "steps", ["id"])

    # Add new columns
    op.add_column("steps", sa.Column("origin", sa.String(), nullable=True))
    op.add_column("steps", sa.Column("organization_id", sa.String(), nullable=True))
    op.add_column("steps", sa.Column("provider_id", sa.String(), nullable=True))
    op.add_column("steps", sa.Column("provider_name", sa.String(), nullable=True))
    op.add_column("steps", sa.Column("model", sa.String(), nullable=True))
    op.add_column("steps", sa.Column("context_window_limit", sa.Integer(), nullable=True))
    op.add_column(
        "steps",
        sa.Column("completion_tokens_details", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True),
    )
    op.add_column(
        "steps",
        sa.Column("tags", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True),
    )
    op.add_column("steps", sa.Column("tid", sa.String(), nullable=True))

    # Add new foreign key constraint for provider_id
    op.create_foreign_key("fk_steps_organization_id", "steps", "providers", ["provider_id"], ["id"], ondelete="RESTRICT")

    # Add new foreign key constraint for provider_id
    op.create_foreign_key("fk_steps_provider_id", "steps", "organizations", ["organization_id"], ["id"], ondelete="RESTRICT")

    # Add new foreign key constraint for provider_id
    op.create_foreign_key("fk_steps_job_id", "steps", "jobs", ["job_id"], ["id"], ondelete="SET NULL")

    # Drop old step_id and step_count columns which aren't in the new model
    op.drop_column("steps", "step_id")
    op.drop_column("steps", "step_count")

    # Add step_id to messages table
    op.add_column("messages", sa.Column("step_id", sa.String(), nullable=True))
    op.create_foreign_key("fk_messages_step_id", "messages", "steps", ["step_id"], ["id"], ondelete="SET NULL")
    # ### end Alembic commands ###


def downgrade() -> None:
    # Skip this migration for SQLite
    if not settings.letta_pg_uri_no_default:
        return

    # ### commands auto generated by Alembic - please adjust! ###
    # Remove step_id from messages first to avoid foreign key conflicts
    op.drop_constraint("fk_messages_step_id", "messages", type_="foreignkey")
    op.drop_column("messages", "step_id")

    # Restore old step_count and step_id column
    op.add_column("steps", sa.Column("step_count", sa.Integer(), nullable=True))
    op.add_column("steps", sa.Column("step_id", sa.String(), nullable=True))

    # Drop new columns and constraints
    op.drop_constraint("fk_steps_provider_id", "steps", type_="foreignkey")
    op.drop_constraint("fk_steps_organization_id", "steps", type_="foreignkey")
    op.drop_constraint("fk_steps_job_id", "steps", type_="foreignkey")

    op.drop_column("steps", "tid")
    op.drop_column("steps", "tags")
    op.drop_column("steps", "completion_tokens_details")
    op.drop_column("steps", "context_window_limit")
    op.drop_column("steps", "model")
    op.drop_column("steps", "provider_name")
    op.drop_column("steps", "provider_id")
    op.drop_column("steps", "organization_id")
    op.drop_column("steps", "origin")

    # Add constraints back
    op.execute("DELETE FROM steps WHERE job_id IS NULL")
    op.alter_column("steps", "job_id", nullable=False)
    op.create_foreign_key("fk_job_usage_statistics_job_id", "steps", "jobs", ["job_id"], ["id"], ondelete="CASCADE")

    # Change id field from string back to int
    op.add_column("steps", sa.Column("old_id", sa.Integer(), nullable=True))
    op.execute("""UPDATE steps SET old_id = CAST(ABS(hashtext(REPLACE(id, 'step-', '')::text)) AS integer)""")
    op.drop_column("steps", "id")
    op.execute("ALTER TABLE steps RENAME COLUMN old_id TO id")
    op.alter_column("steps", "id", nullable=False)
    op.create_primary_key("pk_steps_id", "steps", ["id"])

    # Rename the table
    op.rename_table("steps", "job_usage_statistics")
    # ### end Alembic commands ###