#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""This module allows to connect to a MySQL database."""
import json
from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union

from airflow.hooks.dbapi import DbApiHook
from airflow.models import Connection

if TYPE_CHECKING:
    from mysql.connector.abstracts import MySQLConnectionAbstract
    from MySQLdb.connections import Connection as MySQLdbConnection

MySQLConnectionTypes = Union['MySQLdbConnection', 'MySQLConnectionAbstract']


class MySqlHook(DbApiHook):
    """
    Interact with MySQL.

    You can specify charset in the extra field of your connection
    as ``{"charset": "utf8"}``. Also you can choose cursor as
    ``{"cursor": "SSCursor"}``. Refer to the MySQLdb.cursors for more details.

    Note: For AWS IAM authentication, use iam in the extra connection parameters
    and set it to true. Leave the password field empty. This will use the
    "aws_default" connection to get the temporary token unless you override
    in extras.
    extras example: ``{"iam":true, "aws_conn_id":"my_aws_conn"}``

    :param schema: The MySQL database schema to connect to.
    :type schema: Optional[str]
    :param connection: The :ref:`MySQL connection id <howto/connection:mysql>` used for MySQL credentials.
    :type connection: Optional[Dict]
    """

    conn_name_attr = 'mysql_conn_id'
    default_conn_name = 'mysql_default'
    conn_type = 'mysql'
    hook_name = 'MySQL'
    supports_autocommit = True

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.schema = kwargs.pop("schema", None)
        self.connection = kwargs.pop("connection", None)

    def set_autocommit(self, conn: MySQLConnectionTypes, autocommit: bool) -> None:
        """
        The MySQLdb (mysqlclient) client uses an `autocommit` method rather
        than an `autocommit` property to set the autocommit setting

        :param conn: connection to set autocommit setting
        :type MySQLConnectionTypes: connection object.
        :param autocommit: autocommit setting
        :type bool: True to enable autocommit, False to disable autocommit
        :rtype: None
        """
        if hasattr(conn.__class__, 'autocommit') and isinstance(conn.__class__.autocommit, property):
            conn.autocommit = autocommit
        else:
            conn.autocommit(autocommit)

    def get_autocommit(self, conn: MySQLConnectionTypes) -> bool:
        """
        The MySQLdb (mysqlclient) client uses a `get_autocommit` method
        rather than an `autocommit` property to get the autocommit setting

        :param conn: connection to get autocommit setting from.
        :type MySQLConnectionTypes: connection object.
        :return: connection autocommit setting
        :rtype: bool
        """
        if hasattr(conn.__class__, 'autocommit') and isinstance(conn.__class__.autocommit, property):
            return conn.autocommit
        else:
            return conn.get_autocommit()

    def _get_conn_config_mysql_client(self, conn: Connection) -> Dict:
        conn_config = {