"""Snowflake SSO session wrapper with cached connections and refresh retry.

This wrapper standardizes how development code reuses authenticated Snowflake
sessions to avoid repeated browser SSO prompts on every query.
"""
from __future__ import annotations

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

from .snowflake_pool import sf_pool

_AUTH_ERROR_MARKERS = (
    "390112",  # session/token expired
    "390114",
    "250001",
    "authentication token has expired",
    "session no longer exists",
    "session has expired",
)


def _bool_env(name: str, default: bool) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "y", "on"}


class SnowflakeSessionManager:
    """Execute Snowflake queries with cached session reuse and auth refresh."""

    def __init__(
        self,
        refresh_buffer_min: Optional[int] = None,
        max_retries: Optional[int] = None,
        health_check_seconds: int = 60,
        enabled: Optional[bool] = None,
    ) -> None:
        self.refresh_buffer_min = (
            refresh_buffer_min
            if refresh_buffer_min is not None
            else int(os.getenv("SNOWFLAKE_SESSION_REFRESH_BUFFER_MIN", "5"))
        )
        self.max_retries = (
            max_retries
            if max_retries is not None
            else int(os.getenv("SNOWFLAKE_SESSION_MAX_RETRIES", "1"))
        )
        self.health_check_seconds = max(1, health_check_seconds)
        self.enabled = _bool_env("SNOWFLAKE_SESSION_CACHE_ENABLED", True) if enabled is None else enabled
        self._lock = threading.Lock()
        self._last_health_check: Dict[str, float] = {}

    def _cache_key(
        self,
        connection_name: Optional[str],
        database: str,
        schema: str,
        use_env: bool,
    ) -> str:
        conn = connection_name or "default"
        scope = "env" if use_env else "named"
        return f"{scope}|{conn}|{database}|{schema}"

    @staticmethod
    def _is_auth_error(exc: Exception) -> bool:
        message = str(exc).lower()
        return any(marker in message for marker in _AUTH_ERROR_MARKERS)

    def _get_connection(
        self,
        connection_name: Optional[str],
        database: str,
        schema: str,
        use_env: bool,
    ) -> Any:
        if use_env:
            return sf_pool.get_from_env(database=database, schema=schema)
        return sf_pool.get(connection_name=connection_name or "default", database=database, schema=schema)

    def _invalidate(
        self,
        connection_name: Optional[str],
        database: str,
        schema: str,
        use_env: bool,
    ) -> None:
        sf_pool.close_key(
            connection_name=connection_name or "default",
            database=database,
            schema=schema,
            from_env=use_env,
        )

    def _should_health_check(self, key: str) -> bool:
        now = time.time()
        last = self._last_health_check.get(key, 0.0)
        if now - last >= self.health_check_seconds:
            self._last_health_check[key] = now
            return True
        return False

    def warm(
        self,
        connection_name: Optional[str] = None,
        database: str = "",
        schema: str = "",
        use_env: bool = False,
    ) -> None:
        if not self.enabled:
            return
        conn = self._get_connection(connection_name, database, schema, use_env)
        cur = conn.cursor()
        try:
            cur.execute("SELECT 1")
        finally:
            cur.close()

    def invalidate(
        self,
        connection_name: Optional[str] = None,
        database: str = "",
        schema: str = "",
        use_env: bool = False,
    ) -> None:
        self._invalidate(connection_name, database, schema, use_env)

    def execute_query(
        self,
        sql: str,
        *,
        connection_name: Optional[str] = None,
        database: str = "",
        schema: str = "",
        limit: int = 1000,
        use_env: bool = False,
    ) -> list[dict]:
        """Execute query with cached session reuse and auth-expiry refresh."""
        key = self._cache_key(connection_name, database, schema, use_env)

        attempts = 1 if not self.enabled else (self.max_retries + 1)
        last_error: Optional[Exception] = None
        for attempt in range(1, attempts + 1):
            try:
                with self._lock:
                    conn = self._get_connection(connection_name, database, schema, use_env)
                    if self._should_health_check(key):
                        ping = conn.cursor()
                        try:
                            ping.execute("SELECT 1")
                        finally:
                            ping.close()
                cur = conn.cursor()
                try:
                    cur.execute(sql)
                    cols = [c[0] for c in cur.description] if cur.description else []
                    rows = cur.fetchmany(limit) if limit > 0 else cur.fetchall()
                finally:
                    cur.close()
                return [dict(zip(cols, r)) for r in rows]
            except Exception as exc:
                last_error = exc
                if not self.enabled:
                    break
                if attempt >= attempts or not self._is_auth_error(exc):
                    break
                with self._lock:
                    self._invalidate(connection_name, database, schema, use_env)
        if last_error is not None:
            raise last_error
        return []


_SESSION_MANAGER: Optional[SnowflakeSessionManager] = None


def get_snowflake_session_manager() -> SnowflakeSessionManager:
    global _SESSION_MANAGER
    if _SESSION_MANAGER is None:
        _SESSION_MANAGER = SnowflakeSessionManager()
    return _SESSION_MANAGER

