import logging
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Callable, Generator, Iterable, Optional, Union

import snowflake.connector
from snowflake.connector.connection import SnowflakeConnection
from snowflake.connector.cursor import DictCursor, SnowflakeCursor
from snowflake.connector.errors import ProgrammingError

logger = logging.getLogger("snowcap")

UNSUPPORTED_FEATURE = 2
INVALID_COLUMN_ERR = 904
SYNTAX_ERROR = 1003
OBJECT_ALREADY_EXISTS_ERR = 2002
DOES_NOT_EXIST_ERR = 2003
INVALID_IDENTIFIER = 2004
OBJECT_DOES_NOT_EXIST_ERR = 2043
ACCESS_CONTROL_ERR = 3001
ALREADY_EXISTS_ERR = 3041  # Not sure this is correct
INVALID_GRANT_ERR = 3042
FEATURE_NOT_ENABLED_ERR = 3078  # Unsure if this is just Replication Groups or not

connection_params = {
    "account": os.environ.get("SNOWFLAKE_ACCOUNT"),
    "user": os.environ.get("SNOWFLAKE_USER"),
    "password": os.environ.get("SNOWFLAKE_PASSWORD"),
}

_EXECUTION_CACHE = {}
_EXECUTION_CACHE_LOCK = threading.Lock()
# Track queries currently being executed to prevent duplicate execution
_PENDING_QUERIES: dict[tuple[str, str], threading.Event] = {}


def reset_cache():
    """
    Reset the SQL execution cache.

    This clears cached query results so subsequent queries will re-execute.
    Note: This does NOT clear ACCOUNT_USAGE caches, which are designed to persist
    across queries for performance. The fallback to SHOW GRANTS handles cases
    where newly created grants aren't in the ACCOUNT_USAGE cache.
    """
    global _EXECUTION_CACHE
    _EXECUTION_CACHE = {}


def execute(
    conn_or_cursor: Union[SnowflakeConnection, SnowflakeCursor],
    sql: str,
    cacheable: bool = False,
    empty_response_codes: Optional[list[int]] = None,
) -> list:
    if isinstance(sql, str):
        sql_text = sql
    else:
        raise Exception(f"Unknown sql type: {type(sql)}, {sql}")

    if isinstance(conn_or_cursor, SnowflakeConnection):
        session = conn_or_cursor
        cur = session.cursor(snowflake.connector.DictCursor)
    elif isinstance(conn_or_cursor, DictCursor):
        # Already a DictCursor, use it directly
        session = conn_or_cursor.connection
        cur = conn_or_cursor
    elif isinstance(conn_or_cursor, SnowflakeCursor):
        session = conn_or_cursor.connection
        cur = conn_or_cursor
        cur._use_dict_result = True
    elif hasattr(conn_or_cursor, "connection") and hasattr(conn_or_cursor, "execute"):
        # Handle cursor-like objects
        session = conn_or_cursor.connection
        cur = conn_or_cursor
        cur._use_dict_result = True
    else:
        # Undocumented snowpark-specific type snowflake.connector.connection.StoredProcConnection
        # raise Exception(f"Unknown connection type: {type(conn_or_cursor)}, {conn_or_cursor}")
        session = conn_or_cursor
        cur = session.cursor(snowflake.connector.DictCursor)

    if sql.startswith("USE ROLE"):
        desired_role = sql.split(" ", 2)[-1]
        # TODO: determine if this works at all for quoted roles
        if desired_role == session.role:
            return [[]]

    session_header = f"[{session.user}:{session.role}] > {sql_text}"

    # Thread-safe cache check with pending query deduplication
    if cacheable:
        cache_key = (session.role, sql_text)
        with _EXECUTION_CACHE_LOCK:
            # Check if result is already cached
            if session.role in _EXECUTION_CACHE and sql_text in _EXECUTION_CACHE[session.role]:
                result = _EXECUTION_CACHE[session.role][sql_text]
                # logger.warning(f"{session_header}    \033[94m({len(result)} rows, cached)\033[0m")
                return result

            # Check if another thread is already executing this query
            if cache_key in _PENDING_QUERIES:
                # Wait for the other thread to complete
                pending_event = _PENDING_QUERIES[cache_key]
                # Release lock while waiting
                _EXECUTION_CACHE_LOCK.release()
                try:
                    pending_event.wait(timeout=600)  # Wait up to 10 minutes
                finally:
                    _EXECUTION_CACHE_LOCK.acquire()
                # After waiting, check cache again
                if session.role in _EXECUTION_CACHE and sql_text in _EXECUTION_CACHE[session.role]:
                    result = _EXECUTION_CACHE[session.role][sql_text]
                    return result
                # If still not cached, fall through to execute

            # Mark this query as pending
            _PENDING_QUERIES[cache_key] = threading.Event()

    # Log start message for potentially slow queries (ACCOUNT_USAGE, large SHOW commands)
    is_slow_query = "ACCOUNT_USAGE" in sql_text or "GRANTS_TO_ROLES" in sql_text
    if is_slow_query:
        logger.info(f"{session_header}    \033[93m(running...)\033[0m")

    start = time.time()
    try:
        cur.execute(sql_text)
        result = cur.fetchall()
        runtime = time.time() - start
        # Log execution details - use debug for routine queries, info for slow ones
        if is_slow_query:
            logger.info(f"{session_header}    \033[94m({len(result)} rows, {runtime:.2f}s)\033[0m")
        else:
            logger.debug(f"{session_header}    \033[94m({len(result)} rows, {runtime:.2f}s)\033[0m")
        if cacheable:
            cache_key = (session.role, sql_text)
            with _EXECUTION_CACHE_LOCK:
                if session.role not in _EXECUTION_CACHE:
                    _EXECUTION_CACHE[session.role] = {}
                _EXECUTION_CACHE[session.role][sql_text] = result
                # Signal waiting threads and cleanup
                if cache_key in _PENDING_QUERIES:
                    _PENDING_QUERIES[cache_key].set()
                    del _PENDING_QUERIES[cache_key]
        return result
    except ProgrammingError as err:
        if empty_response_codes and err.errno in empty_response_codes:
            runtime = time.time() - start
            logger.debug(f"{session_header}    \033[94m(empty, {runtime:.2f}s)\033[0m")
            if cacheable:
                cache_key = (session.role, sql_text)
                with _EXECUTION_CACHE_LOCK:
                    if session.role not in _EXECUTION_CACHE:
                        _EXECUTION_CACHE[session.role] = {}
                    _EXECUTION_CACHE[session.role][sql_text] = []
                    # Signal waiting threads and cleanup
                    if cache_key in _PENDING_QUERIES:
                        _PENDING_QUERIES[cache_key].set()
                        del _PENDING_QUERIES[cache_key]
            return []
        # On error, also cleanup pending query marker
        if cacheable:
            cache_key = (session.role, sql_text)
            with _EXECUTION_CACHE_LOCK:
                if cache_key in _PENDING_QUERIES:
                    _PENDING_QUERIES[cache_key].set()
                    del _PENDING_QUERIES[cache_key]
        logger.error(f"{session_header}    \033[31m(err {err.errno}, {time.time() - start:.2f}s)\033[0m")
        raise ProgrammingError(f"{err} on {sql_text}", errno=err.errno) from err


def execute_in_parallel(
    conn_or_cursor: Union[SnowflakeConnection, SnowflakeCursor],
    sqls: Iterable[tuple[str, Any]],
    error_handler: Optional[Callable[[Exception, str], None]] = None,
    cacheable: bool = False,
    empty_response_codes: Optional[list[int]] = None,
    max_workers: int = 8,
) -> Generator[tuple[Any, list], None, None]:
    """
    Execute a function in parallel using ThreadPoolExecutor and yield results as they complete.

    Args:
        conn_or_cursor: SnowflakeConnection or SnowflakeCursor to use for execution
        sqls: An iterable of SQL statements to execute, optionally with associated items
        error_handler: Optional function to handle errors, takes (Exception, sql) as arguments (
        cacheable: Whether to cache results (default: False)
        empty_response_codes: List of error codes that should return empty results (default: None)
        max_workers: Maximum number of worker threads (default: 8)

    Yields:
        Tuples of (original_sql, result) as they complete

    Example:
        >>> sql_statements = [
        ...     "SELECT COUNT(*) FROM table1",
        ...     "SELECT COUNT(*) FROM table2",
        ...     "SELECT COUNT(*) FROM table3",
        ... ]
        >>> for sql, result in execute_in_parallel(connection, sql_statements):
        ...     print(f"SQL: {sql} produced result: {result}")
    """
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Create a mapping of futures to their original items
        future_to_item = {
            executor.submit(execute, conn_or_cursor, sql, cacheable, empty_response_codes): (sql, item)
            for sql, item in sqls
        }

        # Yield results as they complete
        for future in as_completed(future_to_item):
            sql, item = future_to_item[future]
            try:
                result = future.result()
                yield item, result
            except Exception as e:
                if error_handler:
                    error_handler(e, sql)
                else:
                    logger.error(f"Error processing SQL {sql}: {e}")
                    raise e
