from __future__ import annotations

import base64
import hashlib
import json
import secrets
import threading
import urllib.parse
import webbrowser
import warnings
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Tuple, TYPE_CHECKING

import requests

from ..oauth_callback_server import OAuthCallbackServer
from ..token_cache import TokenCacheFile, mk_key
from ..util import extract_loopback_host, extract_port
from .base import (
    DirectAccessAuthError,
    DirectAccessConfigError,
    DirectAccessOAuthError,
    DirectAccessOAuthInvalidClientError,
    DirectAccessOAuthStateMismatchError,
    DirectAccessOAuthTimeoutError,
    TokenHandler,
    unwrap_secret,
)

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


@dataclass(frozen=True)
class _OAuthTokenErrorContext:
    operation: str
    status: int
    token_url: str
    client_id: str
    has_client_secret: bool
    redirect_uri: str
    requested_role: str | None
    raw_body: str
    payload: dict[str, object] | None

    @property
    def error(self) -> str | None:
        v = (self.payload or {}).get("error")
        return str(v) if isinstance(v, str) and v else None

    @property
    def message(self) -> str:
        p = self.payload or {}
        for key in ("error_description", "message", "detail", "details"):
            v = p.get(key)
            if isinstance(v, str) and v:
                return v
        return self.raw_body

    def details(self) -> dict[str, object]:
        p = self.payload or {}
        return {
            "operation": self.operation,
            "status": int(self.status),
            "error": self.error,
            "message": self.message,
            "code": p.get("code"),
            "token_url": self.token_url,
            "client_id": self.client_id,
            "has_client_secret": bool(self.has_client_secret),
            "redirect_uri": self.redirect_uri,
            "requested_role": self.requested_role,
            "scope": f"session:role:{self.requested_role}" if self.requested_role else None,
            "raw_body": self.raw_body,
        }


def _safe_json(resp: requests.Response) -> dict[str, Any] | None:
    try:
        obj = resp.json()
        return obj if isinstance(obj, dict) else {"value": obj}
    except Exception:
        return None


def _is_keyring_secure(keyring_module: Any) -> bool:
    try:
        backend = keyring_module.get_keyring()
    except Exception:
        return False
    backend_name = str(getattr(backend.__class__, "__name__", "") or "")
    backend_module = str(getattr(backend.__class__, "__module__", "") or "")
    if backend_name == "PlaintextKeyring":
        return False
    if backend_module.startswith("keyring.backends.fail"):
        return False
    return True


def _likely_causes(ctx: _OAuthTokenErrorContext) -> list[str]:
    causes: list[str] = []
    err = (ctx.error or "").lower()
    if err == "invalid_client":
        causes.append("OAuth client id/secret is incorrect, missing, or does not match the Snowflake security integration.")
        if not ctx.has_client_secret:
            causes.append("The OAuth integration may be CONFIDENTIAL, but no client secret is configured.")
    else:
        causes.append("The redirect URI may not match the OAuth integration configuration.")
    if ctx.requested_role:
        causes.append(
            "The requested role may be blocked or not pre-authorized for this OAuth integration "
            "(check PRE_AUTHORIZED_ROLES_LIST / blocked roles)."
        )
    return causes


def _format_oauth_error(ctx: _OAuthTokenErrorContext, *, exc_type: type[DirectAccessAuthError]) -> DirectAccessAuthError:
    pretty_causes = "\n".join(f"- {cause}" for cause in _likely_causes(ctx))
    return exc_type(
        f"""OAuth {ctx.operation} failed.
status={ctx.status} error={ctx.error!r} message={ctx.message}

Likely causes:
{pretty_causes}

details={json.dumps(ctx.details(), indent=2, sort_keys=True)}"""
    )


def _oauth_http_error(
    prefix: str,
    resp: requests.Response,
    *,
    token_url: str,
    client_id: str,
    has_client_secret: bool,
    redirect_uri: str,
    requested_role: str | None,
) -> DirectAccessAuthError:
    status = getattr(resp, "status_code", None)
    req_id = None
    try:
        req_id = resp.headers.get("x-snowflake-request-id") or resp.headers.get("x-request-id")
    except Exception:
        req_id = None

    txt = ""
    try:
        txt = (resp.text or "").strip()
    except Exception:
        txt = ""
    if req_id:
        txt = f"{txt} (request_id={req_id})".strip()
    payload = _safe_json(resp)
    ctx = _OAuthTokenErrorContext(
        operation=prefix.lower().replace(" failed", ""),
        status=int(status or 0),
        token_url=token_url,
        client_id=client_id,
        has_client_secret=has_client_secret,
        redirect_uri=redirect_uri,
        requested_role=requested_role,
        raw_body=txt,
        payload=payload if isinstance(payload, dict) else None,
    )
    exc_type = DirectAccessOAuthInvalidClientError if (ctx.error or "").lower() == "invalid_client" else DirectAccessOAuthError
    return _format_oauth_error(ctx, exc_type=exc_type)


class SnowflakeOAuthTokenHandler(TokenHandler):
    # Cache entries are scoped by Snowflake identity plus OAuth integration details.
    _token_cache: dict[tuple[str, ...], Tuple[str, float]] = {}
    _token_cache_created_at: dict[tuple[str, ...], float] = {}
    _refresh_token_cache: dict[tuple[str, ...], Tuple[str, float]] = {}
    _refresh_token_cache_created_at: dict[tuple[str, ...], float] = {}
    _auth_token_cache: dict[tuple[str, ...], Tuple[Tuple[str, str], float]] = {}
    _authenticate_locks: dict[tuple[str, ...], threading.Lock] = {}
    _authenticate_locks_guard = threading.Lock()

    SERVICE_NAME = "RAI_PYREL_OAUTH"

    def __init__(self, conn: "SnowflakeConnectionBase"):
        super().__init__(conn)
        self._file_cache = TokenCacheFile()
        self.client_id = str(getattr(conn, "oauth_client_id", "") or "")
        if not self.client_id:
            raise DirectAccessConfigError("Missing required field oauth_client_id for oauth_authorization_code auth.")
        secret_obj = getattr(conn, "oauth_client_secret", None)
        self.client_secret = unwrap_secret(secret_obj) or ""
        self.redirect_uri = str(getattr(conn, "oauth_redirect_uri", "") or "")
        if not self.redirect_uri:
            raise DirectAccessConfigError("Missing required field oauth_redirect_uri for oauth_authorization_code auth.")
        self.redirect_host = extract_loopback_host(self.redirect_uri)
        if not self.user:
            raise DirectAccessConfigError("user is required for oauth_authorization_code auth.")

        self.redirect_port = extract_port(self.redirect_uri)
        self.account_url = f"https://{self.account}.snowflakecomputing.com"

        self._callback_server: OAuthCallbackServer | None = None
        self._load_refresh_token_cache()
        self._delete_file_token_cache()

    def get_session_login_token(self) -> str:
        token, expired = self._get_access_token_from_cache()
        if token and not expired:
            return token

        with self._authenticate_lock():
            token, expired = self._get_access_token_from_cache()
            if token and not expired:
                return token

            refresh_token, expired_refresh = self._get_refresh_token_from_cache()
            if refresh_token and not expired_refresh:
                try:
                    return self._refresh_token()
                except DirectAccessAuthError:
                    # Refresh tokens can be invalidated; fall back to interactive auth.
                    pass

            auth_code, code_verifier = self._authenticate()
            return self._exchange_code_for_token(auth_code, code_verifier)

    def get_ingress_token(self, endpoint: str) -> str:
        """Return the same OAuth access token for DA ingress requests.

        Unlike JWT auth, the OAuth flow here does not mint endpoint-scoped ingress
        tokens, so the middleware can reuse the session login token directly.
        """
        _ = endpoint
        return self.get_session_login_token()

    def _get_access_token_from_cache(self) -> Tuple[str, bool]:
        return self._get_token_from_cache(
            SnowflakeOAuthTokenHandler._token_cache,
            created_at_cache=SnowflakeOAuthTokenHandler._token_cache_created_at,
        )

    def _get_refresh_token_from_cache(self) -> Tuple[str, bool]:
        return self._get_token_from_cache(
            SnowflakeOAuthTokenHandler._refresh_token_cache,
            created_at_cache=SnowflakeOAuthTokenHandler._refresh_token_cache_created_at,
        )

    def _get_auth_token_from_cache(self) -> Tuple[str, str, bool]:
        key = self._cache_key()
        if key in SnowflakeOAuthTokenHandler._auth_token_cache:
            (auth_code, code_verifier), expiration_timestamp = SnowflakeOAuthTokenHandler._auth_token_cache[key]
            current_timestamp = datetime.now(timezone.utc).timestamp()
            if expiration_timestamp > current_timestamp:
                return auth_code, code_verifier, False
            return "", "", True
        return "", "", False

    @staticmethod
    def _generate_pkce_challenge() -> Tuple[str, str]:
        code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode("utf-8").rstrip("=")
        code_challenge = (
            base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode("utf-8")).digest()).decode("utf-8").rstrip("=")
        )
        return code_verifier, code_challenge

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

    def _authenticate_lock(self) -> threading.Lock:
        key = self._cache_key()
        with SnowflakeOAuthTokenHandler._authenticate_locks_guard:
            lock = SnowflakeOAuthTokenHandler._authenticate_locks.get(key)
            if lock is None:
                lock = threading.Lock()
                SnowflakeOAuthTokenHandler._authenticate_locks[key] = lock
            return lock

    def _get_authorization_url(self, state: str, code_challenge: str) -> str:
        # Only request a role when it was explicitly configured on the connection.
        # Snowflake OAuth integrations can block roles; omitting scope uses the user's default role.
        params = {
            "client_id": self.client_id,
            "response_type": "code",
            "redirect_uri": self.redirect_uri,
            "state": state,
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
        }
        if self.requested_role:
            params["scope"] = f"session:role:{self.requested_role}"
        return f"{self.account_url}/oauth/authorize?" + urllib.parse.urlencode(params)

    def _token_request_common(self) -> tuple[str, dict[str, str], tuple[str, str] | None]:
        token_url = f"{self.account_url}/oauth/token-request"
        headers = {"Content-Type": "application/x-www-form-urlencoded"}
        auth: tuple[str, str] | None = None
        if self.client_secret:
            auth = (self.client_id, self.client_secret)
        return token_url, headers, auth

    def _post_token_request(self, *, operation: str, data: dict[str, str]) -> tuple[requests.Response, datetime]:
        token_url, headers, auth = self._token_request_common()
        issue_time = datetime.now(timezone.utc)
        try:
            resp = requests.post(token_url, data=data, headers=headers, auth=auth, timeout=30)
        except Exception as e:
            raise DirectAccessAuthError(f"Error during OAuth {operation}: {e}") from e
        if resp.status_code != 200:
            raise _oauth_http_error(
                f"Token {operation} failed",
                resp,
                token_url=token_url,
                client_id=self.client_id,
                has_client_secret=bool(self.client_secret),
                redirect_uri=self.redirect_uri,
                requested_role=self.requested_role,
            )
        return resp, issue_time

    def _authenticate(self) -> Tuple[str, str]:
        auth_code, code_verifier, expired = self._get_auth_token_from_cache()
        if auth_code and code_verifier and not expired:
            return auth_code, code_verifier

        code_verifier, code_challenge = self._generate_pkce_challenge()
        state = secrets.token_urlsafe(32)
        auth_url = self._get_authorization_url(state, code_challenge)

        self._callback_server = OAuthCallbackServer(host=self.redirect_host, port=self.redirect_port)
        try:
            self._callback_server.start()
        except Exception as e:
            raise DirectAccessOAuthError(
                f"Failed to start OAuth callback server on {self.redirect_host}:{self.redirect_port}. "
                "Ensure nothing else is using the port and that your redirect URI matches it."
            ) from e

        try:
            browser_open_warning: str | None = None
            try:
                opened = webbrowser.open(auth_url)
            except Exception:
                browser_open_warning = (
                    "Failed to open a browser for OAuth authentication. "
                    f"Open this URL manually:\n{auth_url}"
                )
                warnings.warn(browser_open_warning, UserWarning, stacklevel=2)
            else:
                if opened is False:
                    browser_open_warning = (
                        f"Open this URL in a browser to continue OAuth authentication:\n{auth_url}"
                    )
                    warnings.warn(browser_open_warning, UserWarning, stacklevel=2)

            try:
                result = self._callback_server.wait_for_callback(timeout_s=300.0)
            except TimeoutError as e:
                msg = (
                    "Timed out waiting for OAuth callback. "
                    f"Make sure your Snowflake OAuth integration redirects to {self.redirect_uri!r}."
                )
                if browser_open_warning:
                    msg = f"{msg}\n\n{browser_open_warning}"
                raise DirectAccessOAuthTimeoutError(msg) from e
        finally:
            try:
                self._callback_server.close()
            except Exception:
                pass

        err = str(result.get("error") or "")
        if err:
            raise DirectAccessOAuthError(f"OAuth authentication failed: {err}")
        if str(result.get("state") or "") != state:
            raise DirectAccessOAuthStateMismatchError("OAuth state mismatch - possible security issue.")
        auth_code = str(result.get("code") or "")
        if not auth_code:
            raise DirectAccessOAuthError("OAuth callback did not include an authorization code.")

        expiration_timestamp = (datetime.now(timezone.utc) + timedelta(seconds=60)).timestamp()
        SnowflakeOAuthTokenHandler._auth_token_cache[self._cache_key()] = (
            (auth_code, code_verifier),
            expiration_timestamp,
        )
        return auth_code, code_verifier

    def _exchange_code_for_token(self, auth_code: str, code_verifier: str) -> str:
        SnowflakeOAuthTokenHandler._auth_token_cache.pop(self._cache_key(), None)
        data = {
            "grant_type": "authorization_code",
            "code": auth_code,
            "redirect_uri": self.redirect_uri,
            "code_verifier": code_verifier,
            "client_id": self.client_id,
        }
        resp, issue_time = self._post_token_request(operation="exchange", data=data)
        return self._handle_token_response(resp, issue_time)

    def _refresh_token(self) -> str:
        refresh_token, expired = self._get_refresh_token_from_cache()
        if not refresh_token or expired:
            auth_code, code_verifier = self._authenticate()
            return self._exchange_code_for_token(auth_code, code_verifier)

        data = {"grant_type": "refresh_token", "refresh_token": refresh_token, "client_id": self.client_id}
        resp, issue_time = self._post_token_request(operation="refresh", data=data)
        return self._handle_token_response(resp, issue_time, is_refresh=True)

    def _handle_token_response(self, resp: requests.Response, issue_time: datetime, *, is_refresh: bool = False) -> str:
        data = resp.json()
        key = self._cache_key()
        access_token = str(data.get("access_token") or "")
        expires_in = float(data.get("expires_in") or 0)
        if not access_token or not expires_in:
            raise DirectAccessAuthError("Snowflake OAuth token response missing access_token/expires_in.")
        access_expiration_ts = (issue_time + timedelta(seconds=expires_in)).timestamp()
        access_cached_at_ts = issue_time.timestamp()
        SnowflakeOAuthTokenHandler._token_cache[key] = (access_token, access_expiration_ts)
        SnowflakeOAuthTokenHandler._token_cache_created_at[key] = access_cached_at_ts
        self._delete_file_token_cache()

        refresh_entry: Tuple[str, float] | None = None
        refresh_cached_at_ts: float | None = None
        refresh_token = data.get("refresh_token")
        refresh_expires_in = data.get("refresh_token_expires_in")
        if refresh_token and refresh_expires_in:
            refresh_expiration_ts = (issue_time + timedelta(seconds=float(refresh_expires_in))).timestamp()
            refresh_entry = (str(refresh_token), float(refresh_expiration_ts))
            SnowflakeOAuthTokenHandler._refresh_token_cache[key] = refresh_entry
            refresh_cached_at_ts = issue_time.timestamp()
            SnowflakeOAuthTokenHandler._refresh_token_cache_created_at[key] = refresh_cached_at_ts
        elif is_refresh:
            # Some refresh responses do not rotate the refresh token. Keep the
            # existing entry persisted so a later process restart can still use it.
            refresh_entry = SnowflakeOAuthTokenHandler._refresh_token_cache.get(key)
            refresh_cached_at_ts = SnowflakeOAuthTokenHandler._refresh_token_cache_created_at.get(key)

        if refresh_entry is not None:
            self._persist_refresh_token_cache_entry(refresh_entry, cached_at_ts=refresh_cached_at_ts)
            self._delete_file_token_cache()

        return access_token

    # ---- optional keyring support (best-effort) ----
    def _load_refresh_token_cache(self) -> None:
        try:
            import keyring  # type: ignore
        except Exception:
            return
        try:
            if not _is_keyring_secure(keyring):
                warnings.warn("Insecure keyring detected. Please use a secure keyring backend.", UserWarning, stacklevel=2)
                return
            key = self._cache_key()
            raw = keyring.get_password(SnowflakeOAuthTokenHandler.SERVICE_NAME, json.dumps(key))
            if not raw:
                return
            parsed = json.loads(raw)
            if isinstance(parsed, dict):
                token = parsed.get("token")
                exp = parsed.get("exp")
                if isinstance(token, str) and token and isinstance(exp, (int, float)):
                    SnowflakeOAuthTokenHandler._refresh_token_cache[key] = (token, float(exp))
                    cached_at = parsed.get("cached_at")
                    if isinstance(cached_at, (int, float)):
                        SnowflakeOAuthTokenHandler._refresh_token_cache_created_at[key] = float(cached_at)
                    else:
                        SnowflakeOAuthTokenHandler._refresh_token_cache_created_at.pop(key, None)
                    return
            SnowflakeOAuthTokenHandler._refresh_token_cache[key] = tuple(parsed)  # type: ignore[assignment]
            SnowflakeOAuthTokenHandler._refresh_token_cache_created_at.pop(key, None)
        except Exception:
            warnings.warn("Failed to load oauth refresh token from keyring.", UserWarning, stacklevel=2)

    def _persist_refresh_token_cache_entry(self, cache_entry: Tuple[str, float], *, cached_at_ts: float | None) -> None:
        try:
            import keyring  # type: ignore
        except Exception:
            return
        try:
            if not _is_keyring_secure(keyring):
                warnings.warn("Insecure keyring detected. Please use a secure keyring backend.", UserWarning, stacklevel=2)
                return
            key = self._cache_key()
            keyring.set_password(
                SnowflakeOAuthTokenHandler.SERVICE_NAME,
                json.dumps(key),
                json.dumps(
                    {
                        "token": cache_entry[0],
                        "exp": float(cache_entry[1]),
                        "cached_at": cached_at_ts,
                    }
                ),
            )
        except Exception:
            warnings.warn("Failed to persist oauth refresh token to keyring.", UserWarning, stacklevel=2)

    # ---- legacy file cache cleanup (best-effort) ----
    def _delete_file_token_cache(self) -> None:
        access_key = mk_key(**self._cache_file_key(namespace="oauth_access"))
        refresh_key = mk_key(**self._cache_file_key(namespace="oauth_refresh"))
        self._file_cache.delete(access_key)
        self._file_cache.delete(refresh_key)


__all__ = ["SnowflakeOAuthTokenHandler"]

