"""remove workout_segment.segment_id and workout.gpx

Revision ID: e63433a1d62e
Revises: 141d3978536a
Create Date: 2025-10-15 12:20:54.297530

"""
from alembic import op
import sqlalchemy as sa
import sys
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = 'e63433a1d62e'
down_revision = '141d3978536a'
branch_labels = None
depends_on = None


def upgrade():
    # Ensure all workouts are updated before executing upgrade
    connection = op.get_bind()
    workouts_to_update, segments_to_update = connection.execute(
        sa.text("""SELECT
        (SELECT COUNT(*) 
         FROM workouts 
         WHERE workouts.gpx IS NOT NULL AND workouts.start_point_geom IS NULL),
        (SELECT COUNT(*)
         FROM workout_segments
         WHERE workout_segments.geom IS NULL)""")
    ).one()
    if workouts_to_update or segments_to_update:
        print(
            "\nError:\n"
            "Can not run migration, workouts with file and without geometry "
            "exist.\n"
            "Please downgrade to the previous version of FitTrackee (<1.1.0) "
            "and run Workout Refresh CLI with '--add-missing-geometry' option.\n"
        )
        sys.exit(1)

    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('workout_segments', schema=None) as batch_op:
        batch_op.alter_column('start_date',
               existing_type=postgresql.TIMESTAMP(),
               nullable=False)
        batch_op.drop_column('segment_id')

    with op.batch_alter_table('workouts', schema=None) as batch_op:
        batch_op.drop_column('gpx')
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('workouts', schema=None) as batch_op:
        batch_op.add_column(sa.Column('gpx', sa.VARCHAR(length=255), autoincrement=False, nullable=True))

    with op.batch_alter_table('workout_segments', schema=None) as batch_op:
        batch_op.add_column(sa.Column('segment_id', sa.INTEGER(), autoincrement=False, nullable=True))
        batch_op.alter_column('start_date',
               existing_type=postgresql.TIMESTAMP(),
               nullable=True)

    # ### end Alembic commands ###

    op.execute(
        """
        UPDATE workouts
        SET gpx = original_file
        WHERE original_file like '%.gpx';
        """
    )

    for extension in ['fit', 'kml', 'kmz', 'tcx']:
        op.execute(
            f"""
            UPDATE workouts
            SET gpx = REPLACE (original_file, '.{extension}', '.gpx')
            WHERE original_file like '%.{extension}';
            """
        )

    op.execute(
        """
        UPDATE workout_segments
        SET segment_id = ws.calculated_segment_id
        FROM (
                SELECT workout_id,
                       (ROW_NUMBER() OVER (PARTITION BY workout_id ORDER BY start_date) - 1) AS calculated_segment_id
                FROM workout_segments
                ORDER BY workout_id, start_date
              ) ws
        """
    )
    op.alter_column('workout_segments', 'segment_id', nullable=False)
