import atexit
from typing import Any

from snowflake import connector

from src import config, logger
from src.clients.base_client import BaseClient
from src.constants.queries.snowflake import COPY_INTO
from src.exceptions import SnowflakeError, SnowflakeNotFinishedError


class SnowflakeClient(BaseClient):
    def __init__(self):
        self._connection = connector.connect(
            user=config.Snowflake.USER,
            password=config.Snowflake.PASSWORD,
            account=config.Snowflake.ACCOUNT,
            region=config.Snowflake.REGION,
            warehouse=config.Snowflake.WAREHOUSE,
            database=config.Snowflake.PUBLIC_DATA_MAIN_DATABASE,
            schema=config.Snowflake.PUBLIC_DATA_MAIN_PLAYLISTS_SCHEMA,
            role=config.Snowflake.ROLE,
            authenticator=config.Snowflake.AUTHENTICATOR,
        )
        self._cursor = self._connection.cursor()
        atexit.register(self.cleanup)

    @property
    def is_closed(self):
        return not self._connection or self._connection.is_closed

    def cleanup(self):
        self._cursor.close()
        self._cursor = None
        self._connection.close()
        self._connection = None

    def start_warehouse(self):
        try:
            self._cursor.execute(f"ALTER WAREHOUSE {config.Snowflake.WAREHOUSE} RESUME;")
        except Exception as ex:
            logger.log.warning(f"Cannot start warehouse: {ex}")
        self._cursor.execute(f"USE WAREHOUSE {config.Snowflake.WAREHOUSE};")

    def stop_warehouse(self):
        self._cursor.execute(f"ALTER WAREHOUSE {config.Snowflake.WAREHOUSE} SUSPEND;")

    def get_single_value(self, query: str) -> Any:
        self.start_warehouse()
        return self._cursor.execute(query).fetchone()[0]

    def execute_query_async(self, query: str, s3_path: str) -> str:
        self.start_warehouse()
        query = COPY_INTO.format(s3_path=s3_path, query=query, integration=config.Snowflake.INTEGRATION)
        self._cursor.execute_async(query)
        return self._cursor.sfqid

    def cancel_query(self, query_id: str):
        try:
            self._cursor.execute("SELECT SYSTEM$CANCEL_QUERY(%s)", (query_id,))
        except Exception as ex:
            logger.log.warning(f"Cannot cancel query: {ex}")

    def check_query_finished(self, query_id: str):
        status = self._connection.get_query_status(query_id)
        if self._connection.is_an_error(status):
            raise SnowflakeError(status)
        if self._connection.is_still_running(status):
            raise SnowflakeNotFinishedError(str(status))