# type: ignore
"""load initial version data

Revision ID: 84b534bbc858
Revises: 1e5554a164e6
Create Date: 2025-05-28 17:25:06.159790

"""

import enum
import logging
from datetime import datetime

import sqlalchemy as sa
from alembic import op

revision = "84b534bbc858"
down_revision = "1e5554a164e6"
branch_labels = None
depends_on = None


class Operation(int, enum.Enum):
    INSERT = 0
    UPDATE = 1
    DELETE = 2


logger = logging.getLogger(__name__)

initial_transaction_id = 1
tables = [
    ("iamc_datapoint_universal", "iamc_datapoint_universal_version"),
    ("iamc_measurand", "iamc_measurand_version"),
    ("iamc_timeseries", "iamc_timeseries_version"),
    ("iamc_variable", "iamc_variable_version"),
    ("model", "model_version"),
    ("region", "region_version"),
    ("run", "run_version"),
    ("runmetaentry", "runmetaentry_version"),
    ("scenario", "scenario_version"),
    ("unit", "unit_version"),
]

tabledata_chunksize = 10000


def generate_initial_version_data(
    data_table_name: str, version_table_name: str, transaction_id: int
):
    """Generate initial version data for a given data table."""
    logging.info(
        f"Copying initial data from '{data_table_name}' "
        f"to '{version_table_name}' table ..."
    )

    conn = op.get_bind()

    if conn is not None:
        version_table = sa.Table(
            version_table_name,
            sa.MetaData(),
            autoload_with=conn,
        )

        data_table = sa.Table(
            data_table_name,
            sa.MetaData(),
            autoload_with=conn,
        )

        partitions = (
            conn.execute(
                sa.select(data_table).execution_options(yield_per=tabledata_chunksize)
            )
            .mappings()
            .partitions(tabledata_chunksize)
        )
        for res in partitions:
            initial_versions = [
                {
                    **row,
                    "transaction_id": transaction_id,
                    "end_transaction_id": None,
                    "operation_type": Operation.INSERT,
                }
                for row in res
            ]
            op.bulk_insert(version_table, initial_versions)
    else:
        logging.warning("No database connection available.")


def upgrade():
    logging.info(
        f"Creating initial transaction object with ID {initial_transaction_id} ..."
    )

    conn = op.get_bind()

    if conn is not None:
        if conn.dialect.name != "postgresql":
            logging.info("Not running on postgres database, skipping migration.")
            return

        # reflect transaction table
        transaction_table = sa.Table(
            "transaction",
            sa.MetaData(),
            autoload_with=conn,
        )

        # Check if any transactions already exist
        res = conn.execute(sa.select(transaction_table)).all()
        if len(res) > 0:
            logging.warning(
                "Database already contains transactions. "
                "Skipping creation of initial version "
                "data to avoid corrupting the database."
            )
            return

        # Create initial transaction
        op.bulk_insert(
            transaction_table,
            [
                {
                    "id": initial_transaction_id,
                    "remote_addr": "",
                    "issued_at": datetime.now(),
                },
            ],
        )

    else:
        logging.warning("No database connection available.")

    for data_table_name, version_table_name in tables:
        generate_initial_version_data(
            data_table_name,
            version_table_name,
            initial_transaction_id,
        )


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###
