from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from datetime import datetime, timedelta, timezone
from typing import Any
from typing import TYPE_CHECKING, Tuple

from relationalai.runtime.execution_summary import SUMMARY_EVENT_TOKEN_CACHE_HIT

if TYPE_CHECKING:
    from relationalai.config.connections.snowflake import SnowflakeConnectionBase

_EXEC_LOG = logging.getLogger("relationalai.client.execution")


class DirectAccessAuthError(RuntimeError):
    """Raised for runtime failures during token minting/exchange."""


class DirectAccessConfigError(ValueError):
    """Raised when required auth configuration is missing or invalid."""


class DirectAccessOAuthError(DirectAccessAuthError):
    """Base error for oauth_authorization_code direct access authentication."""


class DirectAccessOAuthInvalidClientError(DirectAccessOAuthError):
    """OAuth token endpoint rejected client credentials."""


class DirectAccessOAuthTimeoutError(DirectAccessOAuthError):
    """OAuth browser login did not redirect back within timeout."""


class DirectAccessOAuthStateMismatchError(DirectAccessOAuthError):
    """OAuth callback state mismatch."""


class DirectAccessOAuthBrowserOpenError(DirectAccessOAuthError):
    """Failed to open a browser window for OAuth login."""


class DirectAccessEndpointStaleError(DirectAccessAuthError):
    """The cached Direct Access endpoint is stale and should be refreshed."""

    def __init__(self, *, status: int, message: str, details: dict[str, object] | None = None):
        super().__init__(message)
        self.status = int(status)
        self.message = str(message)
        self.details = dict(details or {})


def unwrap_secret(v: object) -> str | None:
    try:
        from pydantic import SecretStr  # local import

        if isinstance(v, SecretStr):
            return v.get_secret_value()
    except Exception:
        pass
    if isinstance(v, str):
        return v
    return None


class TokenHandler(ABC):
    """Upstream-like token handler for Direct Access auth.

    A handler is constructed from a single Snowflake connection entry, so all auth
    config lives under `connections.*` (no `direct_access` config group).
    """

    def __init__(self, conn: "SnowflakeConnectionBase"):
        super().__init__()
        self._conn = conn
        raw_role = conn.role
        self.requested_role: str | None = str(raw_role).strip() if isinstance(raw_role, str) and raw_role.strip() else None
        # Preserve the effective Snowflake session role separately from the cache/scope role.
        self.role: str = self.requested_role or "PUBLIC"
        # `user` is not part of SnowflakeConnectionBase, and some authenticators
        # (for example plain OAuth token auth) do not define it at all.
        self.user: str = str(getattr(conn, "user", None) or "")
        self.account: str = self._prepare_account_name(conn.account)
        self._renewal_buffer_time = timedelta(minutes=2)

    @staticmethod
    def from_snowflake_connection(conn: "SnowflakeConnectionBase") -> "TokenHandler":
        auth = str(getattr(conn, "authenticator", "") or "")
        if auth == "oauth_authorization_code":
            from .oauth_authorization_code import SnowflakeOAuthTokenHandler

            return SnowflakeOAuthTokenHandler(conn)
        if auth == "programmatic_access_token":
            from .programmatic_access_token import SnowflakePATHandler

            return SnowflakePATHandler(conn)
        if auth == "jwt":
            from .snowflake_jwt import SnowflakeKeyFileHandler

            return SnowflakeKeyFileHandler(conn)
        raise DirectAccessConfigError(
            f"Unsupported direct access authenticator {auth!r}. "
            "Expected one of: jwt, oauth_authorization_code, programmatic_access_token."
        )

    @abstractmethod
    def get_session_login_token(self) -> str:
        """Return a token suitable for Snowflake session authentication."""

    @abstractmethod
    def get_ingress_token(self, endpoint: str) -> str:
        """Return a token suitable for Direct Access requests to a specific endpoint."""

    @staticmethod
    def _prepare_account_name(account: str) -> str:
        # Replace '_' with '-' to avoid SSL mismatch errors for auth endpoints.
        return account.replace("_", "-").upper()

    def _cache_role_key(self) -> str:
        # Distinguish "no role requested" from explicit PUBLIC.
        return self.requested_role or ""

    def _cache_key(self) -> tuple[str, ...]:
        return (self.account, self.user, self._cache_role_key())

    def _cache_file_key(self, *, namespace: str, extra: str | None = None) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "namespace": namespace,
            "account": self.account,
            "user": self.user,
            "role": self._cache_role_key(),
        }
        if extra is not None:
            payload["extra"] = extra
        return payload

    def _get_token_from_cache(
        self,
        cache: dict[tuple[str, ...], Tuple[str, float]],
        *,
        created_at_cache: dict[tuple[str, ...], float] | None = None,
    ) -> Tuple[str, bool]:
        key = self._cache_key()
        if key in cache:
            token, expiration_timestamp = cache[key]
            current_timestamp = datetime.now(timezone.utc).timestamp()
            if expiration_timestamp - self._renewal_buffer_time.total_seconds() > current_timestamp:
                created_at = created_at_cache.get(key) if created_at_cache is not None else None
                _EXEC_LOG.info(
                    "op=token_cache_hit Using cached Direct Access token for account=%s user=%s role=%s",
                    self.account,
                    self.user,
                    self._cache_role_key() or "<default>",
                    extra={
                        "summary_event": SUMMARY_EVENT_TOKEN_CACHE_HIT,
                        "cache_created_at_ts": created_at,
                    },
                )
                return token, False
            return "", True
        return "", False


__all__ = [
    "TokenHandler",
    "DirectAccessAuthError",
    "DirectAccessConfigError",
    "DirectAccessOAuthError",
    "DirectAccessOAuthInvalidClientError",
    "DirectAccessOAuthTimeoutError",
    "DirectAccessOAuthStateMismatchError",
    "DirectAccessOAuthBrowserOpenError",
    "DirectAccessEndpointStaleError",
    "unwrap_secret",
]

