from typing import Dict, Optional, Set, Tuple, Any
import time
import hashlib
from jose import jwt
from keycloak import KeycloakOpenID
import threading
import urllib3
from pydantic import ConfigDict, BaseModel
from transpara import tlogging

# Disable insecure request warnings
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

logger = tlogging.get_logger(__name__)


class TAuthUser(BaseModel):
    """
    User representation built from Keycloak JWT payload.

    Fields are extracted from standard Keycloak JWT claims:
      - id: 'sub' claim
      - email: 'email' claim
      - first_name: 'given_name' claim
      - surname: 'family_name' claim
      - roles: 'realm_access.roles' claim
      - client_roles: 'resource_access' claim (per-client role sets)
      - raw_claims: full JWT payload for extensibility
    """

    id: str = ""
    email: Optional[str] = None
    first_name: Optional[str] = None
    surname: Optional[str] = None
    roles: Set[str] = set()
    client_roles: Dict[str, Set[str]] = {}
    raw_claims: dict = {}

    model_config = ConfigDict(from_attributes=True)

    @property
    def is_superuser(self) -> bool:
        return "admin" in self.roles

    def has_role(self, role: str) -> bool:
        """Check if user has a specific realm role."""
        return role in self.roles

    def has_any_role(self, *roles: str) -> bool:
        """Check if user has at least one of the specified realm roles."""
        return bool(self.roles & set(roles))

    def has_all_roles(self, *roles: str) -> bool:
        """Check if user has all of the specified realm roles."""
        return set(roles) <= self.roles

    def has_client_role(self, client: str, role: str) -> bool:
        """Check if user has a specific role for a given client."""
        return role in self.client_roles.get(client, set())


class TAuthTokenVerification:
    """
    Token verification and authentication manager for Keycloak user tokens.

    Designed for API services that receive bearer tokens from end users.
    Validates tokens offline using JWKS, periodically checks revocation
    via the introspection endpoint (cached), and builds a TAuthUser from
    the JWT payload without any additional Keycloak round-trips.

    For machine-to-machine (service account) authentication, use TAuthManager instead.

    Features:
      - Offline JWT signature and expiry validation (crypto only, no network)
      - Cached token introspection for revocation detection (configurable TTL)
      - JWKS key caching with automatic refresh
      - Background cache cleanup thread
      - Token forwarding headers for API-to-API chaining
      - Username/password login for OAuth2 password grant flows
    """

    def __init__(
        self,
        server_url: str,
        realm: str,
        client_id: str,
        client_secret: str,
        verify_ssl: bool = True,
        timeout: int = 30,
        jwks_cache_ttl: int = 3600,
        token_cache_ttl: int = 30,
    ):
        """
        Initialize the TAuthTokenVerification instance.

        Args:
            server_url: Keycloak server URL (e.g. "https://auth.example.com/")
            realm: Keycloak realm name
            client_id: Keycloak client ID for this API
            client_secret: Keycloak client secret
            verify_ssl: Whether to verify SSL certificates
            timeout: Connection timeout in seconds
            jwks_cache_ttl: How long to cache JWKS keys (seconds, default 3600)
            token_cache_ttl: How long to cache introspection results (seconds, default 30)
        """
        self.server_url = server_url
        self.realm = realm
        self.client_id = client_id
        self.client_secret = client_secret
        self.verify_ssl = verify_ssl
        self.timeout = timeout
        self.jwks_cache_ttl = jwks_cache_ttl
        self.token_cache_ttl = token_cache_ttl

        # Initialize Keycloak client
        self.keycloak_openid = KeycloakOpenID(
            server_url=server_url,
            client_id=client_id,
            realm_name=realm,
            client_secret_key=client_secret,
            verify=verify_ssl,
            timeout=timeout
        )

        # JWKS cache
        self.jwks_cache = {
            "keys": None,
            "last_updated": 0,
            "lock": threading.Lock()
        }

        # Token introspection cache — key: token SHA-256 hash, value: (is_active, expiry_time)
        self.token_introspection_cache = {}
        self.token_introspection_lock = threading.Lock()

        # Background cleanup thread
        self.cleanup_thread = None
        self.stop_cleanup = threading.Event()

        # Start background cleanup automatically
        self.start_cleanup_scheduler()

    @staticmethod
    def _token_hash(token: str) -> str:
        """Hash a token for use as a cache key to avoid storing full JWTs in memory."""
        return hashlib.sha256(token.encode()).hexdigest()

    def get_jwks(self) -> Dict[str, Any]:
        """Get the JWKS from Keycloak with caching."""
        current_time = time.time()

        with self.jwks_cache["lock"]:
            if self.jwks_cache["keys"] is None or (current_time - self.jwks_cache["last_updated"]) > self.jwks_cache_ttl:
                try:
                    logger.tdebug("Fetching JWKS from Keycloak")
                    jwks = self.keycloak_openid.certs()
                    self.jwks_cache["keys"] = jwks
                    self.jwks_cache["last_updated"] = current_time
                    logger.tdebug("Successfully fetched JWKS from Keycloak")
                except Exception as e:
                    logger.error(f"Error fetching JWKS from Keycloak: {str(e)}")
                    if self.jwks_cache["keys"] is None:
                        logger.error("No cached JWKS available, token validation will fail")
                        raise
                    else:
                        logger.tdebug("Using cached JWKS despite refresh failure")
            return self.jwks_cache["keys"]

    def validate_token_offline(self, token: str) -> Tuple[bool, Dict[str, Any]]:
        """
        Validate a token without contacting Keycloak (signature + expiry only).
        Returns (is_valid, payload).
        """
        try:
            jwks = self.get_jwks()

            options = {
                'verify_signature': True,
                'verify_aud': False,
                'verify_exp': True
            }

            payload = jwt.decode(token, jwks, options=options)
            return True, payload
        except Exception as e:
            logger.debug(f"Token validation failed: {e}")
            return False, {}

    def check_token_revocation(self, token: str) -> bool:
        """
        Check if a token has been revoked by checking the introspection endpoint.
        Results are cached to minimize Keycloak round-trips.
        Returns True if token is active, False if revoked or error.
        """
        current_time = time.time()
        token_key = self._token_hash(token)

        # Check cache first
        with self.token_introspection_lock:
            if token_key in self.token_introspection_cache:
                is_active, expiry_time = self.token_introspection_cache[token_key]
                if current_time < expiry_time:
                    return is_active
                del self.token_introspection_cache[token_key]

        # Not in cache or cache expired, check with Keycloak
        try:
            introspection = self.keycloak_openid.introspect(token)
            is_active = introspection.get('active', False)

            with self.token_introspection_lock:
                self.token_introspection_cache[token_key] = (is_active, current_time + self.token_cache_ttl)

            return is_active
        except Exception as e:
            logger.error(f"Error checking token introspection: {e}")
            return False

    def _build_user(self, payload: Dict[str, Any]) -> TAuthUser:
        """Build a TAuthUser from JWT payload claims — no Keycloak round-trip needed."""
        return TAuthUser(
            id=payload.get('sub', ''),
            email=payload.get('email', ''),
            first_name=payload.get('given_name', ''),
            surname=payload.get('family_name', ''),
            roles=set(payload.get('realm_access', {}).get('roles', [])),
            client_roles={
                client: set(access.get('roles', []))
                for client, access in payload.get('resource_access', {}).items()
            },
            raw_claims=payload,
        )

    def validate_token(self, token: str) -> Tuple[bool, Optional[TAuthUser]]:
        """
        Comprehensive token validation with caching.

        1. Validates JWT signature and expiry offline (crypto only, no network)
        2. Checks revocation via introspection (cached, periodic Keycloak call)
        3. Builds user from JWT payload (no network)

        Returns (is_valid, user_object).
        """
        try:
            is_valid, payload = self.validate_token_offline(token)
            if not is_valid:
                logger.tdebug("Token validation failed: Invalid token structure or signature")
                return False, None

            subject_id = payload.get('sub')
            if not subject_id:
                logger.tdebug("Token validation failed: Missing subject ID in token")
                return False, None

            is_active = self.check_token_revocation(token)
            if not is_active:
                logger.tdebug(f"Token validation failed: Token has been revoked for subject {subject_id}")
                return False, None

            user = self._build_user(payload)
            return True, user
        except Exception as e:
            logger.error(f"Unexpected error during token validation: {str(e)}")
            return False, None

    def login_with_keycloak(self, username: str, password: str) -> Tuple[Optional[str], Optional[int], Optional[TAuthUser]]:
        """
        Authenticate with Keycloak using username/password (resource owner password grant).

        Returns (access_token, expires_in, user) or (None, None, None) on failure.
        """
        try:
            logger.tdebug(f"Attempting Keycloak login for user: {username}")
            token = self.keycloak_openid.token(username, password)

            access_token = token.get('access_token')
            expires_in = token.get('expires_in')

            if not access_token:
                logger.tdebug(f"Keycloak login failed: No access token returned for user {username}")
                return None, None, None

            is_valid, payload = self.validate_token_offline(access_token)
            if not is_valid:
                logger.tdebug(f"Keycloak login failed: Could not decode access token for user {username}")
                return None, None, None

            user = self._build_user(payload)
            logger.tdebug(f"Successfully authenticated user: {username}")

            return access_token, expires_in, user
        except Exception as e:
            logger.error(f"Keycloak authentication error for user {username}: {str(e)}")
            return None, None, None

    # --- Token forwarding for API-to-API chaining ---

    @staticmethod
    def get_forwarding_header(token: str) -> Dict[str, str]:
        """
        Build an Authorization header for forwarding a user's bearer token
        to a downstream API call.

        Use this when API A needs to call API B on behalf of the same user.
        The downstream API validates the same token independently.

        Args:
            token: The bearer token received from the incoming request.

        Returns:
            Dict with the Authorization header ready for httpx/requests.
        """
        return {"Authorization": f"Bearer {token}"}

    # --- Cache management ---

    def cleanup_token_cache(self):
        """Remove expired entries from token introspection cache."""
        current_time = time.time()
        with self.token_introspection_lock:
            expired_tokens = [key for key, (_, expiry_time) in self.token_introspection_cache.items()
                             if current_time > expiry_time]
            for key in expired_tokens:
                del self.token_introspection_cache[key]

    def start_cleanup_scheduler(self):
        """Start the background cache cleanup scheduler."""
        if self.cleanup_thread is not None:
            return

        def cleanup_job():
            while not self.stop_cleanup.is_set():
                try:
                    self.cleanup_token_cache()
                    logger.tdebug("Token introspection cache cleaned up")
                except Exception as e:
                    logger.error(f"Error during cache cleanup: {e}")

                self.stop_cleanup.wait(self.token_cache_ttl)

        self.cleanup_thread = threading.Thread(target=cleanup_job, daemon=True)
        self.cleanup_thread.start()
        logger.info(f"Token cache cleanup scheduled every {self.token_cache_ttl} seconds")

    def stop_cleanup_scheduler(self):
        """Stop the background cache cleanup scheduler."""
        if self.cleanup_thread is not None:
            self.stop_cleanup.set()
            self.cleanup_thread.join(timeout=5)
            self.cleanup_thread = None
            logger.info("Stopped background cache cleanup scheduler")
