#!/usr/bin/env python3

import csv
import logging
import sqlite3
import sys
import subprocess
import os

PRE_CMD_PGSQL = [
    'psql', '-d', 'ong', '-U', 'postgres', '-tA', '-c'
]

PRE_CMD_PGSQL_EXECUTE_FILE = [
    'psql', '-d', 'ong', '-U', 'postgres', '-f'
]

TMP_SQL_FILE = "/tmp/migrate_scim_db_for_domain_change.sql"
NOTIFICATION_DB_PATH = "/usr/syno/etc/notification/notification_db_daemon.sqlite"
LOG_FILE_PATH = "/var/log/sas_domain_migrate.log"

def setup_logger():
    """Setup logger with both console and file output"""
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)

    # Avoid duplicate handlers
    if logger.handlers:
        return logger

    # Console handler
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)

    # File handler
    file_handler = logging.FileHandler(LOG_FILE_PATH)
    file_handler.setLevel(logging.DEBUG)

    # Formatter
    formatter_file = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
    formatter_console = logging.Formatter("%(levelname)s - %(message)s")

    console_handler.setFormatter(formatter_console)
    file_handler.setFormatter(formatter_file)

    logger.addHandler(console_handler)
    logger.addHandler(file_handler)

    return logger

def list_all_uids_in_ong_db():
    logger.info("Migrating Ong database user names...")
    uid_list = []

    sql_select_all_ids = "SELECT external_id FROM scim_id;"
    cmd = PRE_CMD_PGSQL.copy() + [sql_select_all_ids]

    try:
        p_result = subprocess.run(cmd, check=True, text=True, capture_output=True)
        for line in p_result.stdout.splitlines():
            uid_list.append(line.strip())
    except subprocess.CalledProcessError as e:
        logger.error(f"Error executing SQL: {e}")
        raise

    return uid_list

def migrate_ong_db(uid_list, mapping):
    if not uid_list:
        logger.info("No UIDs found in Ong database to migrate.")
        return

    sql_update_scim_user = "BEGIN;\n"
    for old_uid in uid_list:
        if old_uid not in mapping:
            # Means that it is not the domain user we are migrating
            continue
        (new_name, new_uid) = mapping[old_uid]
        sql_update_scim_user += f"""
        DELETE FROM scim_email USING scim_id
        WHERE scim_email.user_internal_id = scim_id.internal_id AND scim_id.external_id = '{new_uid}';

        DELETE FROM scim_photo USING scim_id
        WHERE scim_photo.user_internal_id = scim_id.internal_id AND scim_id.external_id = '{new_uid}';

        DELETE FROM scim_group_members USING scim_id
        WHERE scim_group_members.members = scim_id.internal_id AND scim_id.external_id = '{new_uid}';

        DELETE FROM scim_user USING scim_id
        WHERE scim_user.internal_id = scim_id.internal_id AND scim_id.external_id = '{new_uid}';

        DELETE FROM scim_id WHERE external_id = '{new_uid}';

        UPDATE scim_user
        SET user_name = '{new_name}'
        FROM scim_id
        WHERE scim_user.internal_id = scim_id.internal_id AND scim_id.external_id = '{old_uid}';

        UPDATE scim_id
        SET external_id = '{new_uid}'
        WHERE external_id = '{old_uid}';
        """

    sql_update_scim_user = sql_update_scim_user.replace("\\", "\\\\")
    sql_update_scim_user += "COMMIT;"

    try:
        with open(TMP_SQL_FILE, 'w') as f:
            f.write(sql_update_scim_user)
    except Exception as e:
        logger.error(f"Error writing to temporary SQL file: {e}")
        raise

    cmd = PRE_CMD_PGSQL_EXECUTE_FILE.copy() + [TMP_SQL_FILE]

    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
        logger.info("Ong database migration completed successfully.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Error executing migration SQL in PGSQL: {e.stderr}")
        raise

def list_all_uids_in_notification_db():
    logger.info("Migrating Notification database uids...")
    uid_list = []

    sql_select_all_ids = """
    SELECT uid FROM devices
    UNION
    SELECT uid FROM events;
    """

    try:
        with sqlite3.connect(NOTIFICATION_DB_PATH) as conn:
            cursor = conn.cursor()
            cursor.execute(sql_select_all_ids)
            rows = cursor.fetchall()
            for row in rows:
                uid_list.append(row[0])
    except sqlite3.Error as e:
        logger.error(f"Error reading notification database: {e}")
        raise

    return uid_list

def migrate_notification_db(uid_list, mapping):
    if not uid_list:
        logger.info("No UIDs found in Notification database to migrate.")
        return

    # No need to explicitly wrap BEGIN/COMMIT transaction
    sql_update_uids = ""
    for old_uid in uid_list:
        if str(old_uid) not in mapping:
            # Means that it is not the domain user we are migrating
            continue

        (_, new_uid) = mapping[str(old_uid)]
        sql_update_uids += f"""
        UPDATE devices
        SET uid = {new_uid}
        WHERE uid = {old_uid};

        UPDATE events
        SET uid = {new_uid}
        WHERE uid = {old_uid};
        """

    try:
        with sqlite3.connect(NOTIFICATION_DB_PATH) as conn:
            cursor = conn.cursor()
            cursor.executescript(sql_update_uids)
            conn.commit()
            logger.info("Notification database migration completed successfully.")
    except sqlite3.Error as e:
        logger.error(f"Error updating notification database: {e}")
        raise

def migrate(mapping_file):
    logger.info("Starting domain user migration process...")

    mapping_table = {}
    try:
        with open(mapping_file, 'r', newline='') as csvfile:
            reader = csv.reader(csvfile, delimiter=',')
            next(reader)

            for [_, old_uid, new_name, new_uid] in reader:
                mapping_table[old_uid] = (new_name, new_uid)

    except FileNotFoundError:
        logger.error(f"Error: mapping file '{mapping_file}' not found.")
        raise
    except Exception as e:
        logger.error(f"Error: failed to read mapping file: {e}")
        raise


    try:
        # Migrate Ong Database
        uid_list = list_all_uids_in_ong_db() # uid is string type
        migrate_ong_db(uid_list, mapping_table)

        # Migrate notification SQLite DB
        uid_list = list_all_uids_in_notification_db()
        migrate_notification_db(uid_list, mapping_table)
    except Exception as e:
        logger.error(f"Error during migration process: {e}")
        raise

def is_root():
    return os.geteuid() == 0

def usage():
    """Print usage information"""
    print("Usage:")
    print(f"  {sys.argv[0]} <mapping_file>")

logger = setup_logger()

def main():
    if not is_root():
        print("This script must be run as root.")
        return 1

    if len(sys.argv) != 2:
        usage()
        return 1

    mapping_file = sys.argv[1]
    if not os.path.isfile(mapping_file):
        print(f"Error: Mapping file '{mapping_file}' does not exist.")
        return 1

    try:
        migrate(mapping_file)
        logger.info(f"Domain user migration process completed.")
    except Exception as e:
        logger.error(f"Migration failed: {e}")
        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())