from __future__ import annotations

import logging
import threading
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Tuple

import jwt
import requests

from ..jwt_generator import JWTGenerator
from ..token_cache import TokenCacheFile, mk_key
from ..util import remove_url_scheme
from relationalai.runtime.execution_summary import SUMMARY_EVENT_TOKEN_CACHE_HIT
from relationalai.services._shared.endpoint_stale import is_endpoint_stale_payload

from .base import DirectAccessAuthError, DirectAccessConfigError, DirectAccessEndpointStaleError, TokenHandler, unwrap_secret

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

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


class SnowflakeKeyFileHandler(TokenHandler):
    # Cache exchanged ingress tokens by (account, user, role, endpoint_without_scheme).
    _token_cache: dict[Tuple[str, str, str, str], Tuple[str, float]] = {}
    _token_cache_created_at: dict[Tuple[str, str, str, str], float] = {}
    _lock_registry = threading.Lock()
    _endpoint_locks: dict[Tuple[str, str, str, str], threading.RLock] = {}
    _endpoint_conds: dict[Tuple[str, str, str, str], threading.Condition] = {}
    _endpoint_refreshing: set[Tuple[str, str, str, str]] = set()

    def __init__(self, conn: "SnowflakeConnectionBase"):
        super().__init__(conn)
        self._file_cache = TokenCacheFile()
        self.key_file_path = str(getattr(conn, "private_key_path", "") or "")
        if not self.key_file_path:
            raise DirectAccessConfigError("jwt auth requires `private_key_path` on the connection.")
        pwd_obj = getattr(conn, "private_key_passphrase", None)
        self.key_file_pwd = unwrap_secret(pwd_obj) or ""
        if not self.user:
            raise DirectAccessConfigError("user is required for jwt auth.")

        self._jwt_generator = JWTGenerator(self.account, self.user, self.key_file_path, self.key_file_pwd or None)
        self._load_file_token_cache()

    def get_session_login_token(self) -> str:
        # Generate a fresh short-lived assertion.
        return self._jwt_generator.generate()

    def _cache_key_for_endpoint(self, ep: str) -> Tuple[str, str, str, str]:
        return (self.account, self.user, self._cache_role_key(), str(ep or ""))

    def _get_endpoint_sync_primitives(self, ep: str) -> tuple[threading.RLock, threading.Condition]:
        key = self._cache_key_for_endpoint(ep)
        with SnowflakeKeyFileHandler._lock_registry:
            lock = SnowflakeKeyFileHandler._endpoint_locks.get(key)
            if lock is None:
                lock = threading.RLock()
                SnowflakeKeyFileHandler._endpoint_locks[key] = lock
            cond = SnowflakeKeyFileHandler._endpoint_conds.get(key)
            if cond is None:
                cond = threading.Condition(lock)
                SnowflakeKeyFileHandler._endpoint_conds[key] = cond
            return lock, cond

    def _get_endpoint_lock(self, ep: str) -> threading.RLock:
        lock, _cond = self._get_endpoint_sync_primitives(ep)
        return lock

    def _load_cached_ingress_token(self, ep: str) -> None:
        # First check endpoint-scoped on-disk cache.
        with self._get_endpoint_lock(ep):
            try:
                k = mk_key(**self._cache_file_key(namespace="jwt_ingress", extra=ep))
                obj = self._file_cache.get(k)
                if isinstance(obj, dict):
                    tok = obj.get("token")
                    exp = obj.get("exp")
                    if isinstance(tok, str) and tok and isinstance(exp, (int, float)):
                        key = self._cache_key_for_endpoint(ep)
                        SnowflakeKeyFileHandler._token_cache[key] = (tok, float(exp))
                        cached_at = obj.get("cached_at")
                        if isinstance(cached_at, (int, float)):
                            SnowflakeKeyFileHandler._token_cache_created_at[key] = float(cached_at)
                        else:
                            try:
                                SnowflakeKeyFileHandler._token_cache_created_at[key] = float(self._file_cache.path.stat().st_mtime)
                            except Exception:
                                SnowflakeKeyFileHandler._token_cache_created_at.pop(key, None)
            except Exception:
                pass

    def get_ingress_token(self, endpoint: str) -> str:
        ep = remove_url_scheme(endpoint)
        key = self._cache_key_for_endpoint(ep)
        _lock, cond = self._get_endpoint_sync_primitives(ep)
        with cond:
            self._load_cached_ingress_token(ep)
            token, expired = self._get_token_from_cache_for_endpoint(ep)
            if token and not expired:
                return token
            while key in SnowflakeKeyFileHandler._endpoint_refreshing:
                cond.wait()
                self._load_cached_ingress_token(ep)
                token, expired = self._get_token_from_cache_for_endpoint(ep)
                if token and not expired:
                    return token
            SnowflakeKeyFileHandler._endpoint_refreshing.add(key)

        try:
            assertion = self.get_session_login_token()
            return self._exchange_token(assertion, endpoint)
        finally:
            with cond:
                SnowflakeKeyFileHandler._endpoint_refreshing.discard(key)
                cond.notify_all()

    def _get_token_from_cache_for_endpoint(self, ep: str) -> Tuple[str, bool]:
        with self._get_endpoint_lock(ep):
            key = self._cache_key_for_endpoint(ep)
            if key in SnowflakeKeyFileHandler._token_cache:
                token, expiration_timestamp = SnowflakeKeyFileHandler._token_cache[key]
                current_timestamp = datetime.now(timezone.utc).timestamp()
                if expiration_timestamp - self._renewal_buffer_time.total_seconds() > current_timestamp:
                    created_at = SnowflakeKeyFileHandler._token_cache_created_at.get(key)
                    _EXEC_LOG.info(
                        "op=token_cache_hit Using cached JWT ingress token for account=%s user=%s endpoint=%s",
                        self.account,
                        self.user,
                        ep,
                        extra={
                            "summary_event": SUMMARY_EVENT_TOKEN_CACHE_HIT,
                            "cache_created_at_ts": created_at,
                        },
                    )
                    return token, False
                return "", True
            return "", False

    def _exchange_token(self, assertion: str, endpoint: str) -> str:
        endpoint_wo_scheme = remove_url_scheme(endpoint)
        # Only request a role scope when it was explicitly configured.
        scope_role = f"session:role:{self.requested_role}" if self.requested_role else None
        scope = f"{scope_role} {endpoint_wo_scheme}" if scope_role else endpoint_wo_scheme

        data = {
            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
            "scope": scope,
            "assertion": assertion,
        }
        url = f"https://{self.account}.snowflakecomputing.com/oauth/token"

        try:
            resp = requests.post(url, data=data, timeout=30)
        except Exception as e:
            raise DirectAccessAuthError(f"Error during token exchange: {e}") from e
        if resp.status_code != 200:
            if is_endpoint_stale_payload(getattr(resp, "text", None) or ""):
                raise DirectAccessEndpointStaleError(
                    status=int(resp.status_code),
                    message=f"JWT ingress token exchange failed for stale endpoint {endpoint_wo_scheme}",
                    details={"endpoint": endpoint_wo_scheme, "body": getattr(resp, "text", None) or ""},
                )
            raise DirectAccessAuthError(f"Token exchange failed: {resp.status_code} - {getattr(resp, 'text', '')}")

        token = str(resp.text or "").strip()
        try:
            exp = jwt.decode(token, options={"verify_signature": False}).get("exp")
            exp_ts = float(exp) if exp is not None else (datetime.now(timezone.utc) + timedelta(minutes=10)).timestamp()
        except Exception:
            exp_ts = (datetime.now(timezone.utc) + timedelta(minutes=10)).timestamp()

        cached_at_ts = datetime.now(timezone.utc).timestamp()
        with self._get_endpoint_lock(endpoint_wo_scheme):
            key = self._cache_key_for_endpoint(endpoint_wo_scheme)
            SnowflakeKeyFileHandler._token_cache[key] = (token, float(exp_ts))
            SnowflakeKeyFileHandler._token_cache_created_at[key] = float(cached_at_ts)
        self._persist_file_ingress_token(endpoint, token, float(exp_ts), float(cached_at_ts))
        return token

    def _load_file_token_cache(self) -> None:
        # Endpoint-scoped ingress token is loaded lazily in `get_ingress_token(endpoint)`.
        return

    def _persist_file_ingress_token(self, endpoint: str, token: str, exp_ts: float, cached_at_ts: float) -> None:
        ep = remove_url_scheme(endpoint)
        k = mk_key(**self._cache_file_key(namespace="jwt_ingress", extra=ep))
        self._file_cache.set(k, {"token": token, "exp": float(exp_ts), "cached_at": float(cached_at_ts)})


__all__ = ["SnowflakeKeyFileHandler"]

