# Copyright (c) 2023 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 re
from decimal import Decimal
from enum import Enum
from functools import lru_cache
from typing import Any, Optional, NewType, cast, NamedTuple, Mapping, List

import boto3
import snowflake
from aws_lambda_powertools import Logger, Tracer
from aws_xray_sdk.core import xray_recorder
from botocore.exceptions import ClientError
from cachetools import LRUCache, cached
from cachetools.keys import hashkey
from czc import json
from snowflake.connector import InterfaceError, ProgrammingError
from snowflake.connector.converter import SnowflakeConverter
from snowflake.connector.network import ReauthenticationRequest

from . import constants, openssl_keys

tracer = Tracer(service='snowflake')
logger = Logger(service=constants.PROJECT_NAME)
cache: Any = LRUCache(maxsize=32)

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

DEFAULT_SCHEMA = 'PUBLIC'
ERR_OBJECT_DOES_NOT_EXIST = 2003
ERR_SCHEMA_DOES_NOT_EXIST = 2043
ERR_EXECUTION_CANCELED = 604
ERR_LOCK_WAITERS_EXCEEDS_LIMIT = 625

SnowflakeConnection = NewType('SnowflakeConnection', object)


@lru_cache(maxsize=1)
def _get_sm_client():
    return boto3.client('secretsmanager')


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


class AsyncQueryId(NamedTuple):
    query_id: str


class ClusterNotFoundError(Exception):
    pass


class SnowflakeConnectionError(Exception):
    pass


class UndefinedTableError(Exception):
    pass


class UndefinedSchemaError(Exception):
    pass


class AsyncQueryError(Exception):
    pass


class DatabaseBusyError(Exception):
    pass


def connect(cluster_identifier: str, db_user: str,
            db_name: str = None, purpose: str = None,
            warehouse: str = None, aws_session: Any = None) -> SnowflakeConnection:
    conn = _connect_internal(cluster_identifier, db_user, db_name,
                             purpose, warehouse, aws_session)
    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, aws_session)


def _verify_connection(conn: SnowflakeConnection):
    try:
        execute(conn, 'SELECT 1', timeout=1)
        return True
    except TimeoutError:
        logger.warning('Connection verification: SELECT 1 took longer than 1s.')
        return False
    except SnowflakeConnectionError:
        return False


def _sm_client(aws_session):
    if aws_session is None:
        return _get_sm_client()
    return aws_session.client('secretsmanager')


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

    credentials = json.loads(secret['SecretString'])
    key_pair = openssl_keys.loads(credentials['private_key'])

    return cast(SnowflakeConnection, 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 _add_label_to_sql(sql: str, query_metadata: Optional[dict]):
    xray_segment = xray_recorder.current_segment()
    if query_metadata is None and xray_segment is None:
        return sql

    query_metadata_with_xray = (dict(**(query_metadata or {}), xray_trace_id=xray_segment.trace_id)
                                if xray_segment else query_metadata)

    label = f'/*CZ:DATA:{json.dumps(query_metadata_with_xray)}*/'
    return re.sub(r'(\w+)\s+', f'\\1 {label} ', sql, count=1)


def _execute_sql(executor):
    try:
        return executor()
    except (InterfaceError, ReauthenticationRequest) 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_SCHEMA_DOES_NOT_EXIST:
            raise UndefinedSchemaError() from error
        elif error.errno == ERR_EXECUTION_CANCELED:
            raise TimeoutError() from error
        elif error.errno == ERR_LOCK_WAITERS_EXCEEDS_LIMIT:
            raise DatabaseBusyError() from error
        raise


def execute(conn: SnowflakeConnection, sql: str,
            timeout: int = None,
            metadata: dict = None,
            fetch_result: bool = True) -> 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:
            @tracer.capture_method
            def snowflake_cursor_execute():
                tracer.put_annotation('FUNCTION_NAME', 'snowflake_cursor_execute')
                curs.execute(_add_label_to_sql(sql, metadata), timeout=timeout)
                tracer.put_metadata('query_id', curs.sfqid)

            snowflake_cursor_execute()
            return [{k.lower(): v for k, v in x.items()} for x in curs.fetchall()] if fetch_result else []

    return _execute_sql(executor)


def execute_batch(conn: SnowflakeConnection, sql: str, timeout: int = None, metadata: dict = None) -> 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(_add_label_to_sql(sql, metadata), timeout=timeout, return_cursors=False)
        except Exception:
            snowflake_conn.rollback()
            raise
        snowflake_conn.commit()

    _execute_sql(executor)


def check_async_query(conn: SnowflakeConnection, serialized_query_id: str) -> AsyncQueryStatus:
    snowflake_conn = cast(snowflake.connector.SnowflakeConnection, conn)
    query_id = AsyncQueryId(**json.loads(serialized_query_id))
    query_status = snowflake_conn.get_query_status(query_id.query_id)
    if snowflake_conn.is_an_error(query_status):
        raise AsyncQueryError(f'Async query {query_id.query_id} failed')
    if snowflake_conn.is_still_running(query_status):
        return AsyncQueryStatus.running
    return AsyncQueryStatus.success


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


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 ''}"))
