"""Persistent Snowflake connection pool — authenticate once, reuse everywhere.

Creates one connection per (account, user, database) key and keeps it alive
for the duration of the pipeline run using CLIENT_SESSION_KEEP_ALIVE.
This eliminates repeated SSO browser prompts during multi-phase workflows.

Usage::

    from src.data_modeling.snowflake_pool import sf_pool

    # Get or create a connection (authenticates only on first call)
    conn = sf_pool.get("default", database="ENERTIA_RAW", schema="ENERTIA_DBO")

    # Use it for queries
    cur = conn.cursor()
    cur.execute("SELECT 1")

    # At pipeline end, close all connections
    sf_pool.close_all()
"""
from __future__ import annotations

import logging
import os
import threading
import time
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)

# Session keep-alive parameters applied to every connection
_SESSION_PARAMS = {
    "CLIENT_SESSION_KEEP_ALIVE": True,
    "STATEMENT_TIMEOUT_IN_SECONDS": 1200,
}


class SnowflakeConnectionPool:
    """Thread-safe singleton pool of persistent Snowflake connections.

    Connections are keyed by (connection_name, database, schema).
    The same connection is reused for every call with matching key.
    """

    def __init__(self) -> None:
        self._connections: Dict[str, Any] = {}  # key → connection
        self._lock = threading.Lock()
        self._timestamps: Dict[str, float] = {}  # key → creation time

    def _key(self, connection_name: str, database: str = "", schema: str = "") -> str:
        return f"{connection_name}|{database}|{schema}"

    def get(
        self,
        connection_name: str = "default",
        database: str = "",
        schema: str = "",
    ) -> Any:
        """Get or create a persistent Snowflake connection.

        Args:
            connection_name: Snowflake CLI connection name from config.toml
            database: Optional database override
            schema: Optional schema override

        Returns:
            An open snowflake.connector.Connection
        """
        key = self._key(connection_name, database, schema)

        with self._lock:
            conn = self._connections.get(key)
            if conn is not None:
                # Verify connection is still alive
                try:
                    conn.cursor().execute("SELECT 1")
                    return conn
                except Exception:
                    logger.info("Connection %s stale, reconnecting...", key)
                    try:
                        conn.close()
                    except Exception:
                        pass
                    del self._connections[key]

            # Create new connection
            conn = self._create(connection_name, database, schema)
            self._connections[key] = conn
            self._timestamps[key] = time.time()
            return conn

    def get_from_env(
        self,
        database: str = "",
        schema: str = "",
    ) -> Any:
        """Get or create a connection using SNOWFLAKE_* env vars.

        This is for functions that use env vars instead of connection names.
        """
        key = self._key("__env__", database, schema)

        with self._lock:
            conn = self._connections.get(key)
            if conn is not None:
                try:
                    conn.cursor().execute("SELECT 1")
                    return conn
                except Exception:
                    logger.info("Env connection %s stale, reconnecting...", key)
                    try:
                        conn.close()
                    except Exception:
                        pass
                    del self._connections[key]

            conn = self._create_from_env(database, schema)
            self._connections[key] = conn
            self._timestamps[key] = time.time()
            return conn

    def _create(self, connection_name: str, database: str, schema: str) -> Any:
        """Create a new Snowflake connection via CLI connection name."""
        import snowflake.connector
        from .snowflake_utils import get_snowflake_connection

        logger.info("Creating persistent Snowflake connection: %s (db=%s, schema=%s)",
                     connection_name, database, schema)

        # Try native connection_name first (Snowflake CLI >=3.x)
        try:
            kwargs: Dict[str, Any] = {
                "connection_name": connection_name,
                "client_store_temporary_credential": True,
                "session_parameters": _SESSION_PARAMS.copy(),
            }
            if database:
                kwargs["database"] = database
            if schema:
                kwargs["schema"] = schema
            conn = snowflake.connector.connect(**kwargs)
            logger.info("Connected via connection_name=%s", connection_name)
            return conn
        except Exception as e:
            logger.debug("connection_name failed (%s), falling back to params", e)

        # Fallback: read config.toml manually
        params = get_snowflake_connection(connection_name)
        if not params or not params.get("account"):
            raise RuntimeError(
                f"Snowflake CLI connection not found: {connection_name}"
            )

        kwargs = {
            "account": params.get("account", ""),
            "user": params.get("user", ""),
            "password": params.get("password", ""),
            "role": params.get("role", ""),
            "warehouse": params.get("warehouse", ""),
            "authenticator": params.get("authenticator", ""),
            "client_store_temporary_credential": True,
            "session_parameters": _SESSION_PARAMS.copy(),
        }
        if database:
            kwargs["database"] = database
        if schema:
            kwargs["schema"] = schema

        conn = snowflake.connector.connect(**kwargs)
        logger.info("Connected via CLI params for %s", connection_name)
        return conn

    def _create_from_env(self, database: str, schema: str) -> Any:
        """Create a connection using SNOWFLAKE_* environment variables."""
        import snowflake.connector

        logger.info("Creating persistent Snowflake connection from env (db=%s)", database)

        kwargs: Dict[str, Any] = {
            "account": os.getenv("SNOWFLAKE_ACCOUNT", ""),
            "user": os.getenv("SNOWFLAKE_USER", ""),
            "password": os.getenv("SNOWFLAKE_PASSWORD", ""),
            "role": os.getenv("SNOWFLAKE_ROLE", ""),
            "warehouse": os.getenv("SNOWFLAKE_WAREHOUSE", ""),
            "authenticator": os.getenv("SNOWFLAKE_AUTHENTICATOR", ""),
            "client_store_temporary_credential": True,
            "session_parameters": _SESSION_PARAMS.copy(),
        }
        if database:
            kwargs["database"] = database
        if schema:
            kwargs["schema"] = schema

        return snowflake.connector.connect(**kwargs)

    def close_all(self) -> None:
        """Close all pooled connections. Call at pipeline end."""
        with self._lock:
            for key, conn in self._connections.items():
                try:
                    conn.close()
                    logger.info("Closed connection: %s", key)
                except Exception:
                    pass
            self._connections.clear()
            self._timestamps.clear()

    def close_key(
        self,
        connection_name: str = "default",
        database: str = "",
        schema: str = "",
        from_env: bool = False,
    ) -> None:
        """Close and remove a specific pooled connection."""
        key = self._key("__env__" if from_env else connection_name, database, schema)
        with self._lock:
            conn = self._connections.pop(key, None)
            self._timestamps.pop(key, None)
        if conn is not None:
            try:
                conn.close()
                logger.info("Closed connection: %s", key)
            except Exception:
                pass

    def status(self) -> Dict[str, Any]:
        """Return pool status for debugging."""
        with self._lock:
            now = time.time()
            return {
                "active_connections": len(self._connections),
                "connections": {
                    key: {
                        "age_seconds": round(now - self._timestamps.get(key, now), 1),
                    }
                    for key in self._connections
                },
            }


# Module-level singleton
sf_pool = SnowflakeConnectionPool()
