import json
import logging

from psycopg2.errorcodes import DUPLICATE_COLUMN
from sqlalchemy import *
from sqlalchemy.engine import reflection

from alembic.migration import MigrationContext
from alembic.operations import Operations

from rhodecode.lib.dbmigrate.versions import _reset_base
from rhodecode.lib.jsonalchemy import MutationObj, JsonType
from rhodecode.model import meta, init_model_encryption


def upgrade(migrate_engine):
    """
    Upgrade operations go here.
    Don't create your own engine; bind migrate_engine to your metadata
    """
    _reset_base(migrate_engine)

    from rhodecode.lib.dbmigrate.schema import db_4_20_0_1 as db

    init_model_encryption(db)

    context = MigrationContext.configure(migrate_engine.connect())
    op = Operations(context)
    inspector = inspect(migrate_engine)

    pr_tables = [db.PullRequest.__table__, db.PullRequestVersion.__table__]
    for pr_table in pr_tables:
        existing_columns = [col["name"] for col in inspector.get_columns(pr_table.name)]
        new_column_name = "ai_code_review_state"
        if new_column_name not in existing_columns:
            with op.batch_alter_table(pr_table.name) as batch_op:
                new_column = Column(
                    new_column_name,
                    MutationObj.as_mutable(
                        JsonType(dialect_map=dict(mysql=UnicodeText(16384))),
                    ),
                    default=dict,
                )
                batch_op.add_column(new_column)


def downgrade(migrate_engine):
    pass
