import logging
import os
import sys
from contextlib import contextmanager
from datetime import datetime
from logging.config import fileConfig

from alembic.operations import MigrateOperation, Operations
from sqlalchemy import engine_from_config, event, pool

from alembic import context

# Configure migration logger at INFO level
logger = logging.getLogger("alembic.runtime.migration")
logger.setLevel(logging.INFO)
if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter("[Migration] %(message)s"))
    logger.addHandler(handler)

# Register Snowflake dialect with Alembic (must be done before any migration runs)
try:
    from alembic.ddl.impl import DefaultImpl
    from snowflake.sqlalchemy import URL as SnowflakeURL  # noqa: F401

    class SnowflakeImpl(DefaultImpl):
        """Alembic implementation for Snowflake dialect."""

        __dialect__ = "snowflake"

except ImportError:
    pass  # snowflake-sqlalchemy not installed, skip registration

# Add the project root to sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from backend.sqlite_compat import sqlite3

# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
config = context.config

# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name is not None:
    fileConfig(config.config_file_name)

# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata

# Import your Flask app and db object
from backend.database.base import db
from backend.setup import setup_app
from flask import Flask

# Create a dummy Flask app to get the metadata
app = Flask(__name__)
# Push an application context manually
with app.app_context():
    app_mode = "agenthub"
    if os.environ.get("APP_MODE") == "answers":
        app_mode = "answers"

    setup_app(
        app, db_url=os.environ["DATABASE_URL"], workload_folder_path=os.environ.get("DB_FOLDER_PATH"), dry=True, app_mode=app_mode
    )  # This will initialize db with the app

target_metadata = db.metadata
if os.environ.get("DB_SCHEMA", None):
    target_metadata.schema = os.environ["DB_SCHEMA"]

@contextmanager
def backup_database_on_failure(db_path: str):
    """Create a backup before migration, restore if migration fails.
    This guarantees complete rollback even if SQLite transactions behave
    unexpectedly with batch operations (which can cause implicit commits).
    """
    if not db_path or not os.path.exists(db_path):
        yield
        return

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_path = f"{db_path}.migration_backup.{timestamp}"

    # Create backup using SQLite's backup API
    logger.info(f"Creating backup: {backup_path}")
    source_conn = sqlite3.connect(db_path)
    backup_conn = sqlite3.connect(backup_path)
    try:
        source_conn.backup(backup_conn)
    finally:
        backup_conn.close()
        source_conn.close()
    try:
        yield
        # Success - remove backup
        logger.info("Success, removing backup")
        os.unlink(backup_path)
    except Exception as e:
        # Failure - restore from backup
        logger.error(f"FAILED: {e}")
        logger.error(f"Manual recovery needed from: {backup_path}")
        raise


# SQLAlchemy-based MySQL backup context manager
@contextmanager
def backup_mysql_on_failure(db_url: str):
    """Backup MySQL database using SQLAlchemy (SQL INSERTs) before migration."""
    if not db_url.lower().startswith("mysql"):
        yield
        return
    import re

    from sqlalchemy import MetaData, Table, create_engine, inspect, text

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    db = re.sub(r"^.*\/", "", db_url.split("?")[0])
    backup_file = f"{db}.migration_backup.{timestamp}.sql"
    engine = create_engine(db_url)
    metadata = MetaData()
    metadata.reflect(bind=engine)
    inspector = inspect(engine)
    logger.info(f"Creating MySQL backup (SQL INSERTs): {backup_file}")
    try:
        with open(backup_file, "w", encoding="utf-8") as f:
            for table_name in inspector.get_table_names():
                table = Table(table_name, metadata, autoload_with=engine)
                # Write DELETE FROM for clean restore
                f.write(f"DELETE FROM `{table_name}`;\n")
                # Dump all rows as INSERTs
                with engine.connect() as conn:
                    rows = conn.execute(table.select()).fetchall()
                    if rows:
                        columns = [col.name for col in table.columns]
                        for row in rows:
                            values = []
                            for v in row:
                                if v is None:
                                    values.append("NULL")
                                elif isinstance(v, (int, float)):
                                    values.append(str(v))
                                elif isinstance(v, (bytes, bytearray)):
                                    values.append(f"0x{v.hex()}")
                                else:
                                    # Escape single quotes
                                    s = str(v).replace("'", "''")
                                    values.append(f"'{s}'")
                            f.write(
                                f"INSERT INTO `{table_name}` ({', '.join(f'`{c}`' for c in columns)}) VALUES ({', '.join(values)});\n"
                            )
            f.write("\n")
    except Exception as e:
        logger.warning(f"MySQL SQLAlchemy backup failed: {e}")

    try:
        yield
        logger.info(f"Success, backup at {os.path.abspath(backup_file)} can be deleted if not needed.")
    except Exception as e:
        logger.error(f"FAILED: {e}")
        logger.error("-" * 80)
        logger.error(f"The database migration failed. A backup has been created at:")
        logger.error(f"  {os.path.abspath(backup_file)}")
        logger.error("Recovery options:")

        tables_prefix = os.environ.get("TABLES_PREFIX", "") or ""
        version_table = f"{tables_prefix}alembic_version" if tables_prefix else "alembic_version"

        if inspector.has_table(version_table):
            logger.error("  1. Automatic Downgrade (Recommended):")
            logger.error("     - Run 'alembic downgrade base' to revert all migration changes.")
        else:
            logger.error("  1. Manual Recovery (Required):")
            logger.error("     - The 'alembic_version' table was not found, so an automatic downgrade is not possible.")
            logger.error(f"     - Action: Connect to your MySQL database, drop all tables from '{db}',")
            logger.error(f"       then restore the backup by executing the SQL script: {os.path.abspath(backup_file)}")

        logger.error("-" * 80)
        raise


def run_migrations_offline() -> None:
    """Run migrations in 'offline' mode.

    This configures the context with just a URL
    and not an Engine, though an Engine is acceptable
    here as well.  By skipping the Engine creation
    we don't even need a DBAPI to be available.

    Calls to context.execute() here emit the given string to the
    script output.

    """
    db_url = os.environ["DATABASE_URL"]
    tables_prefix = os.environ.get("TABLES_PREFIX", "") or ""
    version_table = f"{tables_prefix}alembic_version" if tables_prefix else "alembic_version"
    context.configure(
        url=db_url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
        transaction_per_migration=True,
        version_table=version_table,
    )

    with context.begin_transaction():
        context.run_migrations()


def run_migrations_online() -> None:
    """Run migrations in 'online' mode.

    In this scenario we need to create an Engine
    and associate a connection with the context.

    Uses backup_database_on_failure to guarantee complete rollback
    if migration fails (works around SQLite batch operation issues).
    
    Migration runs in 4 automatic steps:
    - Step 1: Add new columns (non-destructive, backward compatible)
    - Step 2: Create new tables (non-destructive, backward compatible)  
    - Step 3: Schema alterations and data migration
    - Step 4: Cleanup old artifacts
    
    If migration fails at any step:
    - Steps 1-2: Simply downgrade plugin version (old code ignores new columns/tables)
    - Steps 3-4: Run 'alembic downgrade base' then downgrade plugin version
    """
    configuration = config.get_section(config.config_ini_section) or {}
    db_url = os.environ["DATABASE_URL"]
    configuration["sqlalchemy.url"] = db_url

    # Extract file path for backup
    db_path = db_url.replace("sqlite:///", "") if db_url.startswith("sqlite:///") else None

    if db_url.lower().startswith("sqlite"):
        backup_ctx = backup_database_on_failure(db_path)
    # elif db_url.lower().startswith("mysql"):
    #     backup_ctx = backup_mysql_on_failure(db_url)
    else:
        backup_ctx = contextmanager(lambda: (yield))()

    connectable = engine_from_config(
        configuration,
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
        echo=True,  # Log SQL operations at INFO level
    )

    if connectable.dialect.name == "sqlite":

        @event.listens_for(connectable, "connect")
        def _sqlite_connect(dbapi_connection, connection_record):
            # Disable sqlite3's implicit transaction handling so SQLAlchemy
            # can emit BEGIN explicitly for transactional DDL.
            dbapi_connection.isolation_level = None
            dbapi_connection.execute("PRAGMA foreign_keys = OFF")


        @event.listens_for(connectable, "begin")
        def _sqlite_begin(conn):
            # Emit our own BEGIN to ensure DDL participates in the transaction.
            conn.exec_driver_sql("BEGIN")


    tables_prefix = os.environ.get("TABLES_PREFIX", "") or ""
    version_table = f"{tables_prefix}alembic_version" if tables_prefix else "alembic_version"

    with connectable.connect() as connection:
        logger.info(f"Version_table: {version_table}")
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
            transactional_ddl=True,
            transaction_per_migration=False,
            version_table=version_table,  # prefixed version table
            version_table_schema=target_metadata.schema,  # needed for PostgreSQL
            include_schemas=True,  # needed for PostgreSQL
        )

        with backup_ctx, context.begin_transaction():
            db_url_lower = db_url.lower()
            schema = os.environ.get("DB_SCHEMA")
            if schema and db_url_lower.startswith("postgresql"):
                from sqlalchemy import text

                context.execute(text("SET search_path TO :schema").bindparams(schema=schema))  # needed for PostgreSQL
            
            context.run_migrations()


if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()
