import pymysql
import os
import sys

def migrate():
    # Source DB (Windows Host) connection details
    # We can use the IP from current DB_HOST or fallback to 172.26.208.1
    src_host = os.getenv('DB_HOST', '172.26.208.1')
    src_user = os.getenv('DB_USER', 'root')
    src_password = os.getenv('DB_PASSWORD', '45278434')
    
    # Target DB (Docker container) connection details
    # Inside container, 'db' refers to the MariaDB service.
    # Note: we connect to target's internal port 3306.
    tgt_host = 'db'
    tgt_user = 'root'
    tgt_password = '45278434'

    databases = ['PM_proto', 'PM_proto_test']

    print(f"Starting migration from Source ({src_host}:3306) to Target ({tgt_host}:3306)...")

    try:
        # Establish connection to source
        src_conn = pymysql.connect(
            host=src_host,
            user=src_user,
            password=src_password,
            charset='utf8mb4'
        )
        print("Connected to Source Database successfully.")
    except Exception as e:
        print(f"Failed to connect to Source Database: {e}")
        sys.exit(1)

    try:
        # Establish connection to target
        tgt_conn = pymysql.connect(
            host=tgt_host,
            user=tgt_user,
            password=tgt_password,
            charset='utf8mb4'
        )
        print("Connected to Target Database successfully.")
    except Exception as e:
        print(f"Failed to connect to Target Database: {e}")
        src_conn.close()
        sys.exit(1)

    try:
        with src_conn.cursor() as src_cur, tgt_conn.cursor() as tgt_cur:
            # Disable foreign key checks to prevent dependency order issues during creation
            tgt_cur.execute("SET FOREIGN_KEY_CHECKS = 0")
            for db_name in databases:
                print(f"\n--- Migrating Database: {db_name} ---")
                
                # Check if database exists in source
                src_cur.execute(f"SHOW DATABASES LIKE '{db_name}'")
                if not src_cur.fetchone():
                    print(f"Database {db_name} does not exist in source. Skipping.")
                    continue

                # Create database on target
                tgt_cur.execute(f"CREATE DATABASE IF NOT EXISTS `{db_name}` CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci")
                print(f"Created database {db_name} on target (or already exists).")

                # Switch to database
                src_cur.execute(f"USE `{db_name}`")
                tgt_cur.execute(f"USE `{db_name}`")

                # Get list of tables
                src_cur.execute("SHOW TABLES")
                tables = [r[0] for r in src_cur.fetchall()]
                print(f"Tables to migrate: {tables}")

                for table in tables:
                    print(f"  Migrating table: {table}...")
                    
                    # Get CREATE TABLE statement
                    src_cur.execute(f"SHOW CREATE TABLE `{table}`")
                    create_stmt = src_cur.fetchone()[1]
                    # Replace MySQL 8.0 specific collation with MariaDB compatible collation
                    create_stmt = create_stmt.replace('utf8mb4_0900_ai_ci', 'utf8mb4_general_ci')

                    # Drop table on target if exists
                    tgt_cur.execute(f"DROP TABLE IF EXISTS `{table}`")
                    
                    # Create table on target
                    tgt_cur.execute(create_stmt)

                    # Get data from source
                    src_cur.execute(f"SELECT * FROM `{table}`")
                    rows = src_cur.fetchall()
                    
                    if not rows:
                        print(f"    Table {table} is empty.")
                        continue

                    # Get columns list
                    src_cur.execute(f"DESCRIBE `{table}`")
                    cols = [r[0] for r in src_cur.fetchall()]
                    
                    # Insert data into target
                    col_names = ", ".join([f"`{c}`" for c in cols])
                    placeholders = ", ".join(["%s"] * len(cols))
                    insert_stmt = f"INSERT INTO `{table}` ({col_names}) VALUES ({placeholders})"
                    
                    tgt_cur.executemany(insert_stmt, rows)
                    tgt_conn.commit()
                    print(f"    Successfully migrated {len(rows)} rows into {table}.")

        # Re-enable foreign key checks after completion
        with tgt_conn.cursor() as tgt_cur:
            tgt_cur.execute("SET FOREIGN_KEY_CHECKS = 1")
        print("\nMigration completed successfully!")
    except Exception as e:
        print(f"Error during migration: {e}")
        tgt_conn.rollback()
    finally:
        src_conn.close()
        tgt_conn.close()

if __name__ == "__main__":
    migrate()