# Copyright (c) 2025 Oracle and/or its affiliates.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License, version 2.0, as
# published by the Free Software Foundation.
#
# This program is designed to work with certain software (including
# but not limited to OpenSSL) that is licensed under separate terms,
# as designated in a particular file or component or in included license
# documentation. The authors of MySQL hereby grant you an
# additional permission to link the program and your derivative works
# with the separately licensed software that they have either included with
# the program or referenced in the documentation.
#
# Without limiting anything contained in the foregoing, this file,
# which is part of MySQL Connector/Python, is also subject to the
# Universal FOSS Exception, version 1.0, a copy of which can be found at
# http://oss.oracle.com/licenses/universal-foss-exception.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU General Public License, version 2.0, for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin St, Fifth Floor, Boston, MA 02110-1301  USA
"""General utilities for AI features in MySQL Connector/Python.

Includes helpers for:
- defensive dict copying
- temporary table lifecycle management
- SQL execution and result conversions
- DataFrame to/from SQL table utilities
- schema/table/column name validation
- array-like to DataFrame conversion
"""

import copy
import json
import random
import re
import string

from contextlib import contextmanager
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union

import numpy as np
import pandas as pd

from mysql.ai.utils.atomic_cursor import atomic_transaction

from mysql.connector.abstracts import MySQLConnectionAbstract
from mysql.connector.cursor import MySQLCursorAbstract
from mysql.connector.types import ParamsSequenceOrDictType

VAR_NAME_SPACE = "mysql_ai"
RANDOM_TABLE_NAME_LENGTH = 32

PD_TO_SQL_DTYPE_MAPPING = {
    "int64": "BIGINT",
    "float64": "DOUBLE",
    "object": "LONGTEXT",
    "bool": "BOOLEAN",
    "datetime64[ns]": "DATETIME",
}

DEFAULT_SCHEMA = "mysql_ai"

# Misc Utilities


def copy_dict(options: Optional[dict]) -> dict:
    """
    Make a defensive copy of a dictionary, or return an empty dict if None.

    Args:
        options: param dict or None

    Returns:
        dict
    """
    if options is None:
        return {}

    return copy.deepcopy(options)


@contextmanager
def temporary_sql_tables(
    db_connection: MySQLConnectionAbstract,
) -> Iterator[list[tuple[str, str]]]:
    """
    Context manager to track and automatically clean up temporary SQL tables.

    Args:
        db_connection: Database connection object used to create and delete tables.

    Returns:
        None

    Raises:
        DatabaseError:
            If a database connection issue occurs.
            If an operational error occurs during execution.

    Yields:
        temporary_tables: List of (schema_name, table_name) tuples created during the
            context. All tables in this list are deleted on context exit.
    """
    temporary_tables: List[Tuple[str, str]] = []
    try:
        yield temporary_tables
    finally:
        with atomic_transaction(db_connection) as cursor:
            for schema_name, table_name in temporary_tables:
                delete_sql_table(cursor, schema_name, table_name)


def execute_sql(
    cursor: MySQLCursorAbstract, query: str, params: ParamsSequenceOrDictType = None
) -> None:
    """
    Execute an SQL query with optional parameters using the given cursor.

    Args:
        cursor: MySQLCursorAbstract object to execute the query.
        query: SQL query string to execute.
        params: Optional sequence or dict providing parameters for the query.

    Raises:
        DatabaseError:
            If the provided SQL query/params are invalid
            If the query is valid but the sql raises as an exception
            If a database connection issue occurs.
            If an operational error occurs during execution.

    Returns:
        None
    """
    cursor.execute(query, params or ())


def _get_name() -> str:
    """
    Generate a random uppercase string of fixed length for table names.

    Returns:
        Random string of length RANDOM_TABLE_NAME_LENGTH.
    """
    char_set = string.ascii_uppercase
    return "".join(random.choices(char_set, k=RANDOM_TABLE_NAME_LENGTH))


def get_random_name(condition: Callable[[str], bool], max_calls: int = 100) -> str:
    """
    Generate a random string name that satisfies a given condition.

    Args:
        condition: Callable that takes a generated name and returns True if it is valid.
        max_calls: Maximum number of attempts before giving up (default 100).

    Returns:
        A random string that fulfills the provided condition.

    Raises:
        RuntimeError: If the maximum number of attempts is reached without success.
    """
    for _ in range(max_calls):
        if condition(name := _get_name()):
            return name
    # condition never met
    raise RuntimeError("Reached max tries without successfully finding a unique name")


# Format conversions


def format_value_sql(value: Any) -> Tuple[str, List[Any]]:
    """
    Convert a Python value into its SQL-compatible string representation and parameters.

    Args:
        value: The value to format.

    Returns:
        Tuple containing:
            - A string for substitution into a SQL query.
            - A list of parameters to be bound into the query.
    """
    if isinstance(value, (dict, list)):
        if len(value) == 0:
            return "%s", [None]
        return "CAST(%s as JSON)", [json.dumps(value)]
    return "%s", [value]


def sql_response_to_df(cursor: MySQLCursorAbstract) -> pd.DataFrame:
    """
    Convert the results of a cursor's last executed query to a pandas DataFrame.

    Args:
        cursor: MySQLCursorAbstract with a completed query.

    Returns:
        DataFrame with data from the cursor.

    Raises:
        DatabaseError:
            If a database connection issue occurs.
            If an operational error occurs during execution.
            If a compatible SELECT query wasn't the last statement ran
    """

    def _json_processor(elem: Optional[str]) -> Optional[dict]:
        return json.loads(elem) if elem is not None else None

    def _default_processor(elem: Any) -> Any:
        return elem

    idx_to_processor = {}
    for idx, col in enumerate(cursor.description):
        if col[1] == 245:
            # 245 is the MySQL type code for JSON
            idx_to_processor[idx] = _json_processor
        else:
            idx_to_processor[idx] = _default_processor

    rows = cursor.fetchall()

    # Process results
    processed_rows = []
    for row in rows:
        processed_row = list(row)

        for idx, elem in enumerate(row):
            processed_row[idx] = idx_to_processor[idx](elem)

        processed_rows.append(processed_row)

    return pd.DataFrame(processed_rows, columns=cursor.column_names)


def sql_table_to_df(
    cursor: MySQLCursorAbstract, schema_name: str, table_name: str
) -> pd.DataFrame:
    """
    Load the entire contents of a SQL table into a pandas DataFrame.

    Args:
        cursor: MySQLCursorAbstract to execute the query.
        schema_name: Name of the schema containing the table.
        table_name: Name of the table to fetch.

    Returns:
        DataFrame containing all rows from the specified table.

    Raises:
        DatabaseError:
            If the table does not exist
            If a database connection issue occurs.
            If an operational error occurs during execution.
        ValueError: If the schema or table name is not valid
    """
    validate_name(schema_name)
    validate_name(table_name)

    execute_sql(cursor, f"SELECT * FROM {schema_name}.{table_name}")
    return sql_response_to_df(cursor)


# Table operations


def table_exists(
    cursor: MySQLCursorAbstract, schema_name: str, table_name: str
) -> bool:
    """
    Check whether a table exists in a specific schema.

    Args:
        cursor: MySQLCursorAbstract object to execute the query.
        schema_name: Name of the database schema.
        table_name: Name of the table.

    Returns:
        True if the table exists, False otherwise.

    Raises:
        DatabaseError:
            If a database connection issue occurs.
            If an operational error occurs during execution.
        ValueError: If the schema or table name is not valid
    """
    validate_name(schema_name)
    validate_name(table_name)

    cursor.execute(
        """
        SELECT 1
        FROM information_schema.tables
        WHERE table_schema = %s AND table_name = %s
        LIMIT 1
        """,
        (schema_name, table_name),
    )
    return cursor.fetchone() is not None


def delete_sql_table(
    cursor: MySQLCursorAbstract, schema_name: str, table_name: str
) -> None:
    """
    Drop a table from the SQL database if it exists.

    Args:
        cursor: MySQLCursorAbstract to execute the drop command.
        schema_name: Name of the schema.
        table_name: Name of the table to delete.

    Returns:
        None

    Raises:
        DatabaseError:
            If a database connection issue occurs.
            If an operational error occurs during execution.
        ValueError: If the schema or table name is not valid
    """
    validate_name(schema_name)
    validate_name(table_name)

    execute_sql(cursor, f"DROP TABLE IF EXISTS {schema_name}.{table_name}")


def extend_sql_table(
    cursor: MySQLCursorAbstract,
    schema_name: str,
    table_name: str,
    df: pd.DataFrame,
    col_name_to_placeholder_string: Dict[str, str] = None,
) -> None:
    """
    Insert all rows from a pandas DataFrame into an existing SQL table.

    Args:
        cursor: MySQLCursorAbstract for execution.
        schema_name: Name of the database schema.
        table_name: Table to insert new rows into.
        df: DataFrame containing the rows to insert.
        col_name_to_placeholder_string:
            Optional mapping of column names to custom SQL value/placeholder
            strings.

    Returns:
        None

    Raises:
        DatabaseError:
            If the rows could not be inserted into the table, e.g., a type or shape issue
            If a database connection issue occurs.
            If an operational error occurs during execution.
        ValueError: If the schema or table name is not valid
    """
    if col_name_to_placeholder_string is None:
        col_name_to_placeholder_string = {}

    validate_name(schema_name)
    validate_name(table_name)
    for col in df.columns:
        validate_name(str(col))

    qualified_table_name = f"{schema_name}.{table_name}"

    # Iterate over all rows in the DataFrame to build insert statements row by row
    for row in df.values:
        placeholders, params = [], []
        for elem, col in zip(row, df.columns):
            elem = elem.item() if hasattr(elem, "item") else elem

            if col in col_name_to_placeholder_string:
                elem_placeholder, elem_params = col_name_to_placeholder_string[col], [
                    str(elem)
                ]
            else:
                elem_placeholder, elem_params = format_value_sql(elem)

            placeholders.append(elem_placeholder)
            params.extend(elem_params)

        cols_sql = ", ".join([str(col) for col in df.columns])
        placeholders_sql = ", ".join(placeholders)
        insert_sql = (
            f"INSERT INTO {qualified_table_name} "
            f"({cols_sql}) VALUES ({placeholders_sql})"
        )
        execute_sql(cursor, insert_sql, params=params)


def sql_table_from_df(
    cursor: MySQLCursorAbstract, schema_name: str, df: pd.DataFrame
) -> Tuple[str, str]:
    """
    Create a new SQL table with a random name, and populate it with data from a DataFrame.

    If an 'id' column is defined in the dataframe, it will be used as the primary key.

    Args:
        cursor: MySQLCursorAbstract for executing SQL.
        schema_name: Schema in which to create the table.
        df: DataFrame containing the data to be inserted.

    Returns:
        Tuple (qualified_table_name, table_name): The schema-qualified and
        unqualified table names.

    Raises:
        RuntimeError: If a random available table name could not be found.
        ValueError: If any schema, table, or a column name is invalid.
        DatabaseError:
            If a database connection issue occurs.
            If an operational error occurs during execution.
    """
    table_name = get_random_name(
        lambda table_name: not table_exists(cursor, schema_name, table_name)
    )
    qualified_table_name = f"{schema_name}.{table_name}"

    validate_name(schema_name)
    validate_name(table_name)
    for col in df.columns:
        validate_name(str(col))

    columns_sql = []
    for col, dtype in df.dtypes.items():
        # Map pandas dtype to SQL type, fallback is VARCHAR
        sql_type = PD_TO_SQL_DTYPE_MAPPING.get(str(dtype), "LONGTEXT")
        validate_name(str(col))
        columns_sql.append(f"{col} {sql_type}")

    columns_str = ", ".join(columns_sql)

    has_id_col = any(col.lower() == "id" for col in df.columns)
    if has_id_col:
        columns_str += ", PRIMARY KEY (id)"

    # Create table with generated columns
    create_table_sql = f"CREATE TABLE {qualified_table_name} ({columns_str})"
    execute_sql(cursor, create_table_sql)

    try:
        # Insert provided data into new table
        extend_sql_table(cursor, schema_name, table_name, df)
    except Exception:  # pylint: disable=broad-exception-caught
        # Delete table before we lose access to it
        delete_sql_table(cursor, schema_name, table_name)
        raise
    return qualified_table_name, table_name


def validate_name(name: str) -> str:
    """
    Validate that the string is a legal SQL identifier (letters, digits, underscores).

    Args:
        name: Name (schema, table, or column) to validate.

    Returns:
        The validated name.

    Raises:
        ValueError: If the name does not meet format requirements.
    """
    # Accepts only letters, digits, and underscores; change as needed
    if not (isinstance(name, str) and re.match(r"^[A-Za-z0-9_]+$", name)):
        raise ValueError(f"Unsupported name format {name}")

    return name


def source_schema(db_connection: MySQLConnectionAbstract) -> str:
    """
    Retrieve the name of the currently selected schema, or set and ensure the default schema.

    Args:
        db_connection: MySQL connector database connection object.

    Returns:
        Name of the schema (database in use).

    Raises:
        ValueError: If the schema name is not valid
        DatabaseError:
            If a database connection issue occurs.
            If an operational error occurs during execution.
    """
    schema = db_connection.database
    if schema is None:
        schema = DEFAULT_SCHEMA

        with atomic_transaction(db_connection) as cursor:
            create_database_stmt = f"CREATE DATABASE IF NOT EXISTS {schema}"
            execute_sql(cursor, create_database_stmt)

    validate_name(schema)

    return schema


def is_table_empty(
    cursor: MySQLCursorAbstract, schema_name: str, table_name: str
) -> bool:
    """
    Determine if a given SQL table is empty.

    Args:
        cursor: MySQLCursorAbstract with access to the database.
        schema_name: Name of the schema containing the table.
        table_name: Name of the table to check.

    Returns:
        True if the table has no rows, False otherwise.

    Raises:
        DatabaseError:
            If the table does not exist
            If a database connection issue occurs.
            If an operational error occurs during execution.
        ValueError: If the schema or table name is not valid
    """
    validate_name(schema_name)
    validate_name(table_name)

    cursor.execute(f"SELECT 1 FROM {schema_name}.{table_name} LIMIT 1")
    return cursor.fetchone() is None


def convert_to_df(
    arr: Optional[Union[pd.DataFrame, pd.Series, np.ndarray]],
    col_prefix: str = "feature",
) -> Optional[pd.DataFrame]:
    """
    Convert input data to a pandas DataFrame if necessary.

    Args:
        arr: Input data as a pandas DataFrame, NumPy ndarray, pandas Series, or None.

    Returns:
        If the input is None, returns None.
        Otherwise, returns a DataFrame backed by the same underlying data whenever
        possible (except in cases where pandas or NumPy must copy, such as for
        certain views or non-contiguous arrays).

    Notes:
        - If an ndarray is passed, column names will be integer indices (0, 1, ...).
        - If a DataFrame is passed, column names and indices are preserved.
        - The returned DataFrame is a shallow copy and shares data with the original
          input when possible; however, copies may still occur for certain input
          types or memory layouts.
    """
    if arr is None:
        return None

    if isinstance(arr, pd.DataFrame):
        return pd.DataFrame(arr)
    if isinstance(arr, pd.Series):
        return arr.to_frame()

    if arr.ndim == 1:
        arr = arr.reshape(-1, 1)
    col_names = [f"{col_prefix}_{idx}" for idx in range(arr.shape[1])]

    return pd.DataFrame(arr, columns=col_names, copy=False)