# -*- coding: utf-8 -*-
# Copyright (c) 2017 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
import json
import logging
from decimal import Decimal
from enum import Enum
from typing import List, Any, Optional, NewType, cast, NamedTuple, Mapping

import boto3
import snowflake
from botocore.exceptions import ClientError
from cachetools import LRUCache, cached
from retrying import retry
from snowflake.connector import InterfaceError, ProgrammingError
from snowflake.connector.converter import SnowflakeConverter

from czc import constants, openssl_keys

logger = logging.getLogger(constants.PROJECT_NAME)
s3 = boto3.client('s3')
sm = boto3.client('secretsmanager')
cache: Any = LRUCache(maxsize=32)

DictRow = Mapping[str, Any]
QueryResult = List[DictRow]

DEFAULT_SCHEMA = 'PUBLIC'
ERR_OBJECT_DOES_NOT_EXIST = 2003
ERR_LOCK_WAITERS_EXCEEDS_LIMIT = 625


SnowflakeConnection = NewType('SnowflakeConnection', object)


class AsyncQueryStatus(Enum):
    none = 'none'
    running = 'running'
    success = 'success'
    error = 'error'


class AsyncQueryId(NamedTuple):
    session_id: str
    query_id: str


class ClusterNotFoundError(Exception):
    pass


class SnowflakeConnectionError(Exception):
    pass


class UndefinedTableError(Exception):
    pass


class AsyncQueryError(Exception):
    pass


class MissingQueryError(AsyncQueryError):
    pass


class DatabaseBusyError(Exception):
    pass


def connect(cluster_identifier: str, db_user: str, db_name: str = None, purpose: str = None,
            warehouse: str = None) -> SnowflakeConnection:
    conn = _connect_internal(cluster_identifier, db_user, db_name, purpose, warehouse)
    if _verify_connection(conn):
        return conn
    else:
        logger.warning(
            f"Snowflake connection verification failed: {cluster_identifier} db={db_name} user={db_user}.  "
            f"Clear cache and reconnect.")
        cache.clear()
        return _connect_internal(cluster_identifier, db_user, db_name, purpose, warehouse)


def _verify_connection(conn: SnowflakeConnection):
    try:
        execute(conn, 'SELECT 1')
        return True
    except SnowflakeConnectionError:
        return False


@cached(cache)
def _connect_internal(cluster_identifier: str,
                      db_user: str,
                      db_name: Optional[str],
                      purpose: Optional[str],
                      warehouse: Optional[str]) -> SnowflakeConnection:
    logger.debug(f'Get connection for {cluster_identifier} DB:{db_name} USER:{db_user} PURPOSE:{purpose}')
    try:
        secret = sm.get_secret_value(SecretId=f'{db_user}@{cluster_identifier}')
    except ClientError as e:
        if e.response['Error']['Code'] == 'ResourceNotFoundException':
            raise ClusterNotFoundError from e
        raise

    credentials = json.loads(secret['SecretString'])
    key_object = s3.get_object(Bucket=credentials['key_bucket'], Key=credentials['private_key_id'])
    encrypted_private_key = key_object['Body'].read()
    key_pair = openssl_keys.load(encrypted_private_key, credentials['passphrase'])

    return snowflake.connector.connect(database=db_name,
                                       warehouse=warehouse or credentials['default_warehouse'],
                                       user=credentials['user'],
                                       account=credentials['account'],
                                       private_key=key_pair.private_key_bytes,
                                       role=credentials['role'],
                                       schema=DEFAULT_SCHEMA,
                                       autocommit=True)


def _execute_sql(executor):
    try:
        return executor()
    except InterfaceError as error:
        raise SnowflakeConnectionError("Could not connect to database") from error
    except ProgrammingError as error:
        if error.errno == ERR_OBJECT_DOES_NOT_EXIST:
            raise UndefinedTableError() from error
        elif error.errno == ERR_LOCK_WAITERS_EXCEEDS_LIMIT:
            raise DatabaseBusyError() from error
        raise


def execute(conn: SnowflakeConnection, sql: str) -> QueryResult:
    logger.debug(f"Execute Query: [{sql}]")

    def executor():
        snowflake_conn = cast(snowflake.connector.SnowflakeConnection, conn)
        with snowflake_conn.cursor(snowflake.connector.DictCursor) as curs:
            curs.execute(sql)
            return _fetch_result(curs)

    return _execute_sql(executor)


def execute_batch(conn: SnowflakeConnection, sql: str) -> None:
    logger.debug(f"Execute Batch: [{sql}]")

    def executor():
        snowflake_conn = cast(snowflake.connector.SnowflakeConnection, conn)
        snowflake_conn.cursor().execute('BEGIN')
        try:
            snowflake_conn.execute_string(sql, return_cursors=False)
        except Exception:
            snowflake_conn.rollback()
            raise
        snowflake_conn.commit()

    _execute_sql(executor)


# Retry to guard against an initial race: queries do not appear in the query_history views until after compilation
@retry(retry_on_exception=lambda x: isinstance(x, MissingQueryError),
       stop_max_attempt_number=10,
       wait_exponential_multiplier=1000,
       wait_exponential_max=30000)
def check_async_query(conn: SnowflakeConnection, serialized_query_id: str) -> AsyncQueryStatus:
    query_id = AsyncQueryId(**json.loads(serialized_query_id))
    q = f"""
    SELECT execution_status
    FROM table(information_schema.query_history_by_session(
        session_id=>{query_id.session_id},
        result_limit=>10000
    )) WHERE query_id = {quoted_literal(query_id.query_id)}
    """

    result = execute(conn, q)
    if not result:
        raise MissingQueryError(f'Async query {query_id.query_id} not found')

    status = result[0]['execution_status'].lower()
    if status.startswith('failed'):
        raise AsyncQueryError(f'Async query {query_id.query_id} failed')
    elif status == 'success':
        return AsyncQueryStatus.success
    else:
        return AsyncQueryStatus.running


def execute_async(conn: SnowflakeConnection, sql: str) -> str:
    logger.debug(f"Execute Async Query: [{sql}]")
    snowflake_conn = cast(snowflake.connector.SnowflakeConnection, conn)
    with snowflake_conn.cursor() as curs:
        curs.execute(sql, _no_results=True)
        query_info = AsyncQueryId(
            session_id=snowflake_conn.session_id,
            query_id=curs.sfqid
        )
        return json.dumps(query_info._asdict())


def _fetch_result(curs) -> QueryResult:
    try:
        return [{k.lower(): v for k, v in x.items()} for x in curs.fetchall()]
    except ProgrammingError as error:
        if "no results to fetch" not in str(error):
            raise error
    return []


def _canonical_identifier(identifier: str):
    return identifier.replace('"', '""').upper()


def quoted_identifier(*identifiers: str) -> str:
    return '.'.join(f'"{_canonical_identifier(i)}"' for i in identifiers)


def quoted_literal(literal: Any) -> str:
    adjusted_literal = str(literal) if isinstance(literal, Decimal) else literal
    return SnowflakeConverter.quote(SnowflakeConverter.escape(adjusted_literal))


def identifier_literal(identifier: str, prefix: bool = False) -> str:
    return SnowflakeConverter.quote(_canonical_identifier(f"{identifier}{'%' if prefix else ''}"))
