# -*- coding: utf-8 -*-

# BASED ON https://github.com/apache/airflow/blob/349b0811c3022605426ba57d30936240a7c2848a/airflow/providers/snowflake/hooks/snowflake.py
# Original code license
#
# 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.

import logging
import sys

from contextlib import closing
from typing import Any, Iterable, List, Mapping, Tuple, Union

import six

from airflow.contrib.hooks.snowflake_hook import SnowflakeHook


logger = logging.getLogger(__name__)


def snowflake_get_records(
    hook: SnowflakeHook, sql: str, parameters: Union[Mapping, Iterable, None] = None
) -> Tuple[List, int, str]:
    """
    Executes the sql and returns a set of records alongside with
    session_id and query_id generated by Snowflake.

    :param hook: Configured SnowflakeHook with Snoflake connection details.
    :param sql: the sql statement to be executed (str) or a list of
        sql statements to execute
    :param parameters: The parameters to render the SQL query with.
    :return: Query result alongside with session_id and query_id
        assigned by Snowflake to each submitted query
    """
    if sys.version_info[0] < 3:
        sql = sql.encode("utf-8")

    with closing(hook.get_conn()) as conn:
        session_id = conn.session_id
        with closing(conn.cursor()) as cur:
            if parameters is not None:
                cur.execute(sql, parameters)
            else:
                cur.execute(sql)
            query_id = cur.sfqid
            return cur.fetchall(), session_id, query_id


def snowflake_get_first(
    hook: SnowflakeHook, sql: str, parameters: Union[Mapping, Iterable, None] = None
) -> Tuple[Any, int, str]:
    """
    Executes the sql and returns the first resulting row alongside with
    session_id and query_id generated by Snowflake.

    :param hook: Configured SnowflakeHook with Snoflake connection details.
    :param sql: the sql statement to be executed (str) or a list of
        sql statements to execute
    :param parameters: The parameters to render the SQL query with.
    :return: Query result alongside with session_id and query_id
        assigned by Snowflake to each submitted query
    """
    if sys.version_info[0] < 3:
        sql = sql.encode("utf-8")

    with closing(hook.get_conn()) as conn:
        session_id = conn.session_id
        with closing(conn.cursor()) as cur:
            if parameters is not None:
                cur.execute(sql, parameters)
            else:
                cur.execute(sql)
            query_id = cur.sfqid
            return cur.fetchone(), session_id, query_id


def snowflake_run(
    hook: SnowflakeHook,
    sql: Union[str, List[str]],
    autocommit: bool = False,
    parameters: Union[Mapping, Iterable, None] = None,
) -> Tuple[int, List[str]]:
    """
    Runs a command or a list of commands and returns session_id and query_id(s).
    Pass a list of sql statements to the sql parameter to get them to execute sequentially.

    :param hook: Configured SnowflakeHook with Snoflake connection details.
    :param sql: the sql statement to be executed (str) or a list of
        sql statements to execute
    :param autocommit: What to set the connection's autocommit setting to
        before executing the query.
    :param parameters: The parameters to render the SQL query with.
    :return: session_id and query_id(s) assigned by Snowflake to each submitted query
    """
    if isinstance(sql, six.string_types):
        sql = [sql]

    with closing(hook.get_conn()) as conn:
        session_id, query_ids = conn.session_id, []
        if hook.supports_autocommit:
            hook.set_autocommit(conn, autocommit)

        with closing(conn.cursor()) as cur:
            for s in sql:
                if sys.version_info[0] < 3:
                    s = s.encode("utf-8")
                if parameters is not None:
                    hook.log.info("{} with parameters {}".format(s, parameters))
                    res = cur.execute(s, parameters)
                else:
                    hook.log.info(s)
                    res = cur.execute(s)
                query_ids.append(res.sfqid)

        # If autocommit was set to False for db that supports autocommit,
        # or if db does not supports autocommit, we do a manual commit.
        if not hook.get_autocommit(conn):
            conn.commit()
    logger.info(
        "Executed queries '{}', got session_id {}, query_id {}".format(
            sql, session_id, query_ids
        )
    )
    return (session_id, query_ids)


class DbndSnowflakeHook(SnowflakeHook):
    """
    Use this Hook instead of `airflow.contrib.hooks.snowflake_hook.SnowflakeHook` if you
    need session_id and query_id(s) generated by Snowflake.
    E.g. to pass them to `dbnd_snowflake.snowflake_resources.log_snowflake_resource_usage()`.
    """

    @staticmethod
    def from_snowflake_hook(hook: SnowflakeHook) -> "DbndSnowflakeHook":
        return DbndSnowflakeHook(
            snowflake_conn_id=getattr(hook, hook.conn_name_attr),
            warehouse=hook.warehouse,
            account=hook.account,
            database=hook.database,
            region=hook.region,
            role=hook.role,
            schema=hook.schema,
        )

    def get_records(
        self, sql: str, parameters: Union[Mapping, Iterable, None] = None
    ) -> Tuple[List, int, str]:
        """
        Executes the sql and returns a set of records alongside with
        session_id and query_id generated by Snowflake.

        :param sql: the sql statement to be executed (str) or a list of
            sql statements to execute
        :param parameters: The parameters to render the SQL query with.
        :return: Query result alongside with session_id and query_id
            assigned by Snowflake to each submitted query
        """
        return snowflake_get_records(self, sql, parameters)

    def get_first(
        self, sql: str, parameters: Union[Mapping, Iterable, None] = None
    ) -> Tuple[Any, int, str]:
        """
        Executes the sql and returns the first resulting row alongside with
        session_id and query_id generated by Snowflake.

        :param sql: the sql statement to be executed (str) or a list of
            sql statements to execute
        :param parameters: The parameters to render the SQL query with.
        :return: Query result alongside with session_id and query_id
            assigned by Snowflake to each submitted query
        """
        return snowflake_get_first(self, sql, parameters)

    def run(
        self,
        sql: Union[str, List[str]],
        autocommit: bool = False,
        parameters: Union[Mapping, Iterable, None] = None,
    ) -> Tuple[int, List[str]]:
        """
        Runs a command or a list of commands. Pass a list of sql
        statements to the sql parameter to get them to execute
        sequentially

        :param sql: the sql statement to be executed (str) or a list of
            sql statements to execute
        :param autocommit: What to set the connection's autocommit setting to
            before executing the query.
        :param parameters: The parameters to render the SQL query with.
        :return: session_id and query_id(s) assigned by Snowflake to each submitted query
        """
        return snowflake_run(self, sql, autocommit, parameters)
