from logging import WARNING, Logger
from psycopg2 import Binary as PgBinary
from psycopg2.extras import execute_values
from pyodbc import Binary as SqlsBinary
from typing import Any

from .db_common import (
    _db_bind_columns, _db_bind_marks, _db_log, _db_except_msg
)
from .db_pomes import db_connect


def db_migrate_data(errors: list[str],
                    source_engine: str,
                    source_table: str,
                    source_columns: list[str],
                    target_engine: str,
                    target_table: str = None,
                    target_columns: list[str] = None,
                    source_conn: Any = None,
                    target_conn: Any = None,
                    where_clause: tuple = None,
                    batch_size: int = None,
                    logger: Logger = None) -> int:
    """
    Migrate data from one database to another database.

    The origin and destination databases must be in the list of databases configured and
    supported by this package. The columns must not be type LOB (large binary object).
    The columns in the target database must be of equivalent types, respectively,
    in both the origin and the destination database.

    :param errors: incidental error messages
    :param source_engine: the source database engine type
    :param source_table: the origin table
    :param source_columns: the colums to be migrated
    :param target_engine: the destination database engine type
    :param target_table: the table to write the data to (defaults to the source table)
    :param target_columns: the columns to write the data to (defaults to the source columns)
    :param source_conn: the connection to the source database (obtains a new one, if not specified)
    :param target_conn: the connection to the destination database (obtains a new one, if not specified)
    :param where_clause: the criteria for tuple selection
    :param batch_size: the maximum number of tuples in to migrate in each batch, or 0 or None for no limit
    :param logger: optional logger
    :return: number of tuples effectively migrated
    """
    # initialize the return variable
    result: int = 0

    # make sure to have connections to the source and destination databases
    curr_source_conn: Any = source_conn or db_connect(errors=errors,
                                                      engine=source_engine,
                                                      logger=logger)
    curr_target_conn: Any = target_conn or db_connect(errors=errors,
                                                      engine=target_engine,
                                                      logger=logger)

    # make sure to have a target table
    if not target_table:
        target_table = source_table

    # make sure to have target columns
    if not target_columns:
        target_columns = source_columns

    # buid the SELECT query
    cols: str = ", ".join(source_columns)
    sel_stmt: str = f"SELECT {cols} FROM {source_table}"
    if where_clause:
        sel_stmt += f" WHERE {where_clause}"

    # build the INSERT query
    if target_engine == "postgres":
        values: str = "VALUES %s"
    else:
        raw_row: str = _db_bind_marks(engine=target_engine,
                                      start=1,
                                      finish=len(target_columns)+1)
        values: str = f"VALUES({raw_row})"
    cols: str = ", ".join(target_columns)
    insert_stmt = (f"INSERT INTO {target_table} "
                   f"({cols}) {values}")

    # log the migration start
    _db_log(logger=logger,
            engine=source_engine,
            stmt=(f"Started migrating data, "
                  f"from {source_engine}.{source_table} "
                  f"to {target_engine}.{target_table}"))

    # migrate the data
    err_msg: str | None = None
    try:
        source_cursor: Any = curr_source_conn.cursor()
        target_cursor: Any = curr_target_conn.cursor()
        if target_engine == "sqlserver":
            target_cursor.fast_executemany = True

        # execute the query
        source_cursor.execute(statement=sel_stmt)

        # fetch rows in batches/all rows
        rows: list[tuple]
        if batch_size:
            rows = source_cursor.fetchmany(size=batch_size)
        else:
            rows = source_cursor.fetchall()
        while rows:
            match target_engine:
                case "mysql":
                    pass
                case "oracle":
                    target_cursor.executemany(statement=insert_stmt,
                                              parameters=rows)
                    result += len(rows)
                case "postgres":
                    try:
                        execute_values(cur=target_cursor,
                                       sql=insert_stmt,
                                       argslist=rows)
                    except ValueError as e:
                        # is it a 'ValueError' exception on NULLs in strings ?
                        # ("A string literal cannot contain NUL (0x00) characters.")
                        if "contain NUL" in e.args[0]:
                            # yes, log the occurrence, remove the NULLs, and try again
                            _db_log(logger=logger,
                                    engine="postgres",
                                    level=WARNING,
                                    stmt=f"Found NULLs in values for {insert_stmt}")
                            # search for NULLs in input data
                            cleaned_rows: [tuple] = []
                            for raw_row in rows:
                                is_cleaned: bool = False
                                cleaned_row: list = []
                                for val in raw_row:
                                    # is 'val' a string containing NULLs ?
                                    if isinstance(val, str) and val.count(chr(0)) > 0:
                                        # yes, clean it up and mark the row as cleaned
                                        clean_val: str = val.replace(chr(0), "")
                                        is_cleaned = True
                                    else:
                                        clean_val: str = val
                                    cleaned_row.append(clean_val)
                                # has the row been cleaned ?
                                if is_cleaned:
                                    # yes, use it
                                    cleaned_rows.append(tuple(cleaned_row))
                                else:
                                    # no, use original row
                                    cleaned_rows.append(raw_row)
                            # insert the cleaned data
                            execute_values(cur=target_cursor,
                                           sql=insert_stmt,
                                           argslist=cleaned_rows)
                case "sqlserver":
                    target_cursor.executemany(insert_stmt, rows)

            # increment the tuple migration counter
            result += len(rows)
            # read the next batch
            if batch_size:
                rows = source_cursor.fetchmany(size=batch_size)
            else:
                rows = source_cursor.fetchall()

        # close the cursors and commit the transactions
        source_cursor.close()
        curr_source_conn.commit()
        target_cursor.close()
        curr_target_conn.commit()
    except Exception as e:
        # rollback the transactions
        if curr_source_conn:
            curr_source_conn.rollback()
        if curr_target_conn:
            curr_target_conn.rollback()
        err_msg = _db_except_msg(exception=e,
                                 engine=source_engine)
    finally:
        # close the connections, if locally acquired
        if curr_source_conn and not source_conn:
            curr_source_conn.close()
        if curr_target_conn and not target_conn:
            curr_target_conn.close()

    # log the migration finish
    _db_log(logger=logger,
            engine=source_engine,
            err_msg=err_msg,
            errors=errors,
            stmt=(f"Migrated {result} tuples, "
                  f"from {source_engine}.{source_table} "
                  f"to {target_engine}.{target_table}"))

    return result

def db_migrate_lobs(errors: list[str],
                    source_engine: str,
                    source_table: str,
                    source_column: str,
                    pk_columns: list[str],
                    target_engine: str,
                    target_table: str = None,
                    target_column: str = None,
                    source_conn: Any = None,
                    target_conn: Any = None,
                    where_clause: tuple = None,
                    chunk_size: int = None,
                    logger: Logger = None) -> int:
    """
    Migrate large binary objects (LOBs) from one database to another database.

    The origin and destination databases must be in the list of databases configured and
    supported by this package. One or more columns making up a primary key, or a unique row identifier,
    must exist on *source_table* and *target_table*, and be provided in *pk_columns*.
    It is assumed that the primary key columns have the same name, and are of equivalent types,
    in both the origin and the destination database.

    :param errors: incidental error messages
    :param source_engine: the source database engine type
    :param source_table: the table holding the LOBs
    :param source_column: the column holding the LOB
    :param pk_columns: columns making up a primary key, or a unique identifier for a tuple
    :param target_engine: the destination database engine type
    :param target_table: the table to write the lob to (defaults to the source table)
    :param target_column: the column to write the lob to (defaults to the source column)
    :param source_conn: the connection to the source database (obtains a new one, if not specified)
    :param target_conn: the connection to the destination database (obtains a new one, if not specified)
    :param where_clause: the criteria for tuple selection
    :param chunk_size: size in bytes of the data chunk to read/write, or 0 or None for no limit
    :param logger: optional logger
    :return: number of LOBs effectively migrated
    """
    # initialize the return variable
    result: int = 0

    # make sure to have connections to the source and destination databases
    curr_source_conn: Any = source_conn or db_connect(errors=errors,
                                                      engine=source_engine,
                                                      logger=logger)
    curr_target_conn: Any = target_conn or db_connect(errors=errors,
                                                      engine=target_engine,
                                                      logger=logger)

    # make sure to have a target table
    if not target_table:
        target_table = source_table

    # make sure to have a target column
    if not target_column:
        target_column = source_column

    # normalize the chunk size
    if not chunk_size:
        chunk_size = -1

    # buid the SELECT query
    pks: str = ", ".join(pk_columns)
    sel_stmt: str = f"SELECT {pks}, {source_column} FROM {source_table}"
    if where_clause:
        sel_stmt += f" WHERE {where_clause}"
    blob_index: int = len(pk_columns)

    # build the UPDATE query
    set_column: str = _db_bind_columns(engine=target_engine,
                                       columns=[target_column],
                                       concat=", ",
                                       start_index=1)
    where_columns: str = _db_bind_columns(engine=target_engine,
                                          columns=pk_columns,
                                          concat=" AND ",
                                          start_index=2)
    # insert the value
    update_stmt_1: str = (f"UPDATE {target_table} "
                          f"SET {set_column} "
                          f"WHERE {where_columns}")
    # append the value
    update_stmt_2: str = (f"UPDATE {target_table} "
                          f"SET {set_column[:len(target_column)+3]}"
                          f"{target_column} || {set_column[len(target_column)+3:]} "
                          f"WHERE {where_columns}")

    # log the migration start
    _db_log(logger=logger,
            engine=source_engine,
            stmt=(f"Started migrating LOBs, "
                  f"from {source_engine}.{source_table}.{source_column} "
                  f"to {target_engine}.{target_table}.{target_column}"))

    # migrate the LOBs
    err_msg: str | None = None
    try:
        source_cursor: Any = curr_source_conn.cursor()
        target_cursor: Any = curr_target_conn.cursor()

        # execute the query
        source_cursor.execute(statement=sel_stmt)

        # fetch rows
        for row in source_cursor:

            # retrieve the values of the primary key columns (leave blob column out)
            pk_vals: tuple = tuple([row[inx] for inx in range(blob_index)])

            # access the blob in chunks and write it to file
            offset: int = 1
            update_stmt: str = update_stmt_1
            lob: Any = row[blob_index]
            lob_data: bytes = lob.read(offset=offset,
                                       amount=chunk_size)
            while lob_data:
                offset += len(lob_data)
                match target_engine:
                    case "mysql":
                        pass
                    case "oracle":
                        target_cursor.execute(statement=update_stmt,
                                              parameters=(lob_data, *pk_vals))
                    case "postgres":
                        target_cursor.execute(update_stmt, (PgBinary(lob_data), *pk_vals))
                    case "sqlserver":
                        target_cursor.execute(sql=update_stmt,
                                              params=(SqlsBinary(lob_data), *pk_vals))
                update_stmt = update_stmt_2
                # read the next chunk
                lob_data = lob.read(offset=offset,
                                    amount=chunk_size)

            # increment the LOB migration counter
            result += 1

        # close the cursors and commit the transactions
        source_cursor.close()
        curr_source_conn.commit()
        target_cursor.close()
        curr_target_conn.commit()
    except Exception as e:
        # rollback the transactions
        if curr_source_conn:
            curr_source_conn.rollback()
        if curr_target_conn:
            curr_target_conn.rollback()
        err_msg = _db_except_msg(exception=e,
                                 engine=source_engine)
    finally:
        # close the connections, if locally acquired
        if curr_source_conn and not source_conn:
            curr_source_conn.close()
        if curr_target_conn and not target_conn:
            curr_target_conn.close()

    # log the migration finish
    _db_log(logger=logger,
            engine=source_engine,
            err_msg=err_msg,
            errors=errors,
            stmt=(f"Migrated {result} LOBs, "
                  f"from {source_engine}.{source_table}.{source_column} "
                  f"to {target_engine}.{target_table}.{target_column}"))

    return result
