"""Proxy for snowflake requests in order to log them in New Relic."""
from typing import Any, TypeAlias, Mapping

from ddtrace import tracer
import snowflake.connector
from snowflake.connector.connection import SnowflakeConnection
from snowflake_connector import snowflake_conn

from charts.utils.cache import cache_in_redis

DEFAULT_CONFIG = {
    'pool_pre_ping': False,
    'pool_reset_on_return': None,
    'commit_before_close': False
}

SnowflakeParameterValueType: TypeAlias = str | list[str] | list[list[str]] | bool | int
SnowflakeParameterType: TypeAlias = dict[str, SnowflakeParameterValueType]


def pass_rollback(self) -> None:
    """We're only doing SELECT queries. Rollback is not necessary.

    Original implementation in the connector:
        def rollback(self) -> None:
            self.cursor().execute("ROLLBACK")
    """
    pass


SnowflakeConnection.rollback = pass_rollback


@cache_in_redis()
@tracer.wrap(name='snowflake_fetchone')
def fetchone(*args, **kwargs):
    """Passthrough to Snowflake to fetchone."""
    kwargs.update(DEFAULT_CONFIG)
    return snowflake_conn.fetchone(*args, **kwargs)


@cache_in_redis()
@tracer.wrap(name='snowflake_fetchall')
def fetchall(*args, **kwargs):
    """Passthrough to Snowflake to fetchall."""
    kwargs.update(DEFAULT_CONFIG)
    return snowflake_conn.fetchall(*args, **kwargs)


@tracer.wrap(name='snowflake_fetchall')
# TODO: Fix type
def fetchall_nocache(query: str, params: SnowflakeParameterType) -> Any:
    """Passthrough to Snowflake to fetchall without caching."""
    return snowflake_conn.fetchall(query, params, **DEFAULT_CONFIG)


def SQLLoader(path):
    """Create a SQL loader in given path."""
    return snowflake_conn.SQLLoader(path)


# Monkey patch cursor.fetchall
original_fetchall = snowflake.connector.cursor.SnowflakeCursor.fetchall


def traced_fetchall(self, *args, **kwargs):
    with tracer.trace("cursor.fetchall"):
        return original_fetchall(self, *args, **kwargs)


snowflake.connector.cursor.SnowflakeCursor.fetchall = traced_fetchall