from __future__ import annotations

import typing

if typing.TYPE_CHECKING:
    from snowflake.connector import SnowflakeConnection
    from snowflake.connector.cursor import SnowflakeCursor


_connection_cache: dict[str, SnowflakeConnection] = {}


def get_credentials() -> dict[str, str | None]:
    import os
    import json

    raw_credentials = os.environ.get('SNOWFLAKE_CREDENTIALS')
    if raw_credentials is None or raw_credentials == '':
        raise Exception(
            'set credentials as a JSON blob in SNOWFLAKE_CREDENTIALS environment variable'
        )
    credentials: dict[str, str | None] = json.loads(raw_credentials)
    return credentials


def reset_connection(
    credentials: dict[str, str | None] | None = None,
) -> SnowflakeConnection:
    return get_connection(reset=True, credentials=credentials)


def get_connection(
    *,
    reset: bool = False,
    credentials: dict[str, str | None] | None = None,
) -> SnowflakeConnection:
    import json
    import snowflake.connector

    # get serialized credentials
    if credentials is None:
        credentials = get_credentials()
    serialized_credentials = json.dumps(credentials, sort_keys=True)

    if serialized_credentials in _connection_cache and not reset:
        # load from cache if possible
        return _connection_cache[serialized_credentials]
    else:
        # create new connection and cache it with performance optimizations
        performance_params = {
            'client_fetch_use_mp': True,  # Enable multi-processed fetching
            'client_prefetch_threads': 10,  # Max value for maximum parallelism (default: 4, max: 10)
        }

        conn = snowflake.connector.connect(**credentials, **performance_params)

        _connection_cache[serialized_credentials] = conn
        return conn


def get_cursor(
    *,
    conn: SnowflakeConnection | None = None,
    credentials: dict[str, str | None] | None = None,
) -> SnowflakeCursor:
    if conn is None:
        conn = get_connection(credentials=credentials)
    return conn.cursor()


def use_warehouse(
    warehouse: str,
    *,
    conn: SnowflakeConnection | None = None,
    credentials: dict[str, str | None] | None = None,
) -> str:
    cursor = get_cursor(conn=conn, credentials=credentials)
    return cursor.execute('USE WAREHOUSE ' + warehouse).fetchone()[0]  # type: ignore


def use_role(
    role: str,
    *,
    conn: SnowflakeConnection | None = None,
    credentials: dict[str, str | None] | None = None,
) -> None:
    cursor = get_cursor(conn=conn, credentials=credentials)
    return cursor.execute('USE ROLE ' + role).fetchone()[0]  # type: ignore


def use_schema(
    schema: str,
    *,
    conn: SnowflakeConnection | None = None,
    credentials: dict[str, str | None] | None = None,
) -> None:
    cursor = get_cursor(conn=conn, credentials=credentials)
    return cursor.execute('USE SCHEMA ' + schema).fetchone()[0]  # type: ignore


def use_database(
    database: str,
    *,
    conn: SnowflakeConnection | None = None,
    credentials: dict[str, str | None] | None = None,
) -> str:
    cursor = get_cursor(conn=conn, credentials=credentials)
    return cursor.execute('USE DATABASE ' + database).fetchone()[0]  # type: ignore
