backend/services/token_manager.py

import os
import time
import uuid
import json
import jwt
from typing import Optional, Dict, Any, List
from threading import Lock

from backend.services.redis_lock import RedisLockManager

JWT_SECRET = os.getenv("JWT_SECRET", "hyperflow_enterprise_jwt_secret_key_2026")
JWT_ALGORITHM = "HS256"

ACCESS_TOKEN_EXPIRE_MINUTES = int(os.getenv("ACCESS_TOKEN_EXPIRE_MINUTES", "15"))
REFRESH_TOKEN_EXPIRE_DAYS = int(os.getenv("REFRESH_TOKEN_EXPIRE_DAYS", "7"))


class TokenManager:
    """
    Enterprise Dual-Token & JTI Revocation Blacklist Manager.
    Uses Redis for real-time O(1) blacklist lookups with TTL auto-purge.
    Falls back gracefully to a thread-safe in-memory store if Redis is unavailable.
    """

    def __init__(self, redis_client=None):
        self.secret_key = JWT_SECRET
        self.algorithm = JWT_ALGORITHM
        self.lock_manager = RedisLockManager(redis_client=redis_client)
        self.redis = getattr(self.lock_manager, 'client', None)
        self.use_fallback = getattr(self.lock_manager, 'use_fallback', True)

        # In-memory fallbacks when Redis is offline
        self._memory_blacklist: Dict[str, float] = {}  # {jti: exp_timestamp}
        self._memory_user_jtis: Dict[str, List[tuple[str, float]]] = {}  # {sub: [(jti, exp_timestamp)]}
        self._mem_lock = Lock()

    # ---------------------------------------------------------------------------
    # Token Generation
    # ---------------------------------------------------------------------------
    def create_access_token(
        self,
        sub: str,
        role: str = "user",
        scopes: Optional[List[str]] = None,
        custom_jti: Optional[str] = None
    ) -> tuple[str, str, float]:
        """Generates a short-lived access token (15 mins) with a unique JTI."""
        now = time.time()
        exp = now + (ACCESS_TOKEN_EXPIRE_MINUTES * 60)
        jti = custom_jti or str(uuid.uuid4())

        payload = {
            "sub": sub,
            "jti": jti,
            "token_type": "access",
            "role": role,
            "scopes": scopes or ["read"],
            "iat": int(now),
            "exp": int(exp)
        }

        token = jwt.encode(payload, self.secret_key, algorithm=self.algorithm)
        self.track_user_jti(sub, jti, exp)
        return token, jti, exp

    def create_refresh_token(
        self,
        sub: str,
        role: str = "user",
        parent_jti: Optional[str] = None,
        custom_jti: Optional[str] = None
    ) -> tuple[str, str, float]:
        """Generates a long-lived refresh token (7 days) with a unique JTI."""
        now = time.time()
        exp = now + (REFRESH_TOKEN_EXPIRE_DAYS * 86400)
        jti = custom_jti or str(uuid.uuid4())

        payload = {
            "sub": sub,
            "jti": jti,
            "token_type": "refresh",
            "role": role,
            "parent_jti": parent_jti,
            "iat": int(now),
            "exp": int(exp)
        }

        token = jwt.encode(payload, self.secret_key, algorithm=self.algorithm)
        self.track_user_jti(sub, jti, exp)
        return token, jti, exp

    def issue_token_pair(
        self,
        sub: str,
        role: str = "user",
        scopes: Optional[List[str]] = None
    ) -> Dict[str, Any]:
        """Issues an Access + Refresh dual token pair."""
        access_token, access_jti, access_exp = self.create_access_token(sub, role, scopes)
        refresh_token, refresh_jti, refresh_exp = self.create_refresh_token(sub, role, parent_jti=access_jti)

        return {
            "access_token": access_token,
            "refresh_token": refresh_token,
            "token_type": "bearer",
            "expires_in": ACCESS_TOKEN_EXPIRE_MINUTES * 60,
            "access_jti": access_jti,
            "refresh_jti": refresh_jti,
            "access_exp": int(access_exp),
            "refresh_exp": int(refresh_exp)
        }

    # ---------------------------------------------------------------------------
    # Token Verification
    # ---------------------------------------------------------------------------
    def decode_token(self, token_str: str) -> Dict[str, Any]:
        """Decodes and validates token signature and expiration."""
        try:
            payload = jwt.decode(token_str, self.secret_key, algorithms=[self.algorithm])
            return payload
        except jwt.ExpiredSignatureError:
            raise ValueError("Token signature has expired")
        except jwt.InvalidTokenError as e:
            raise ValueError(f"Invalid token: {str(e)}")

    def verify_access_token(self, token_str: str) -> Dict[str, Any]:
        """Decodes access token and ensures it is valid, type=='access', and not blacklisted."""
        payload = self.decode_token(token_str)

        if payload.get("token_type") != "access":
            raise ValueError("Invalid token type. Expected access token.")

        jti = payload.get("jti")
        if not jti:
            raise ValueError("Token payload missing JTI identifier.")

        if self.is_jti_blacklisted(jti):
            raise ValueError("Access token has been revoked (blacklisted JTI).")

        return payload

    def verify_refresh_token(self, token_str: str) -> Dict[str, Any]:
        """Decodes refresh token and ensures type=='refresh' and not blacklisted."""
        payload = self.decode_token(token_str)

        if payload.get("token_type") != "refresh":
            raise ValueError("Invalid token type. Expected refresh token.")

        jti = payload.get("jti")
        if not jti:
            raise ValueError("Token payload missing JTI identifier.")

        if self.is_jti_blacklisted(jti):
            raise ValueError("Refresh token has been revoked (blacklisted JTI).")

        return payload

    # ---------------------------------------------------------------------------
    # Redis JTI Blacklisting Engine
    # ---------------------------------------------------------------------------
    def revoke_jti(self, jti: str, exp_timestamp: float) -> bool:
        """
        Blacklists a JTI in Redis with TTL set to remaining token lifetime.
        Automatically deleted by Redis once token naturally expires.
        """
        now = time.time()
        remaining_ttl = int(exp_timestamp - now)
        if remaining_ttl <= 0:
            return True  # Already naturally expired

        if not self.use_fallback and self.redis:
            try:
                key = f"blacklist:jti:{jti}"
                self.redis.set(key, "revoked", ex=remaining_ttl)
                return True
            except Exception as e:
                print(f"[TokenManager] Redis error during revoke_jti: {e}. Using in-memory fallback.")

        # Fallback to thread-safe in-memory store
        with self._mem_lock:
            self._memory_blacklist[jti] = exp_timestamp
            self._purge_expired_memory_blacklist()
        return True

    def is_jti_blacklisted(self, jti: str) -> bool:
        """Checks whether a JTI exists in the Redis or in-memory blacklist."""
        if not jti:
            return False

        if not self.use_fallback and self.redis:
            try:
                key = f"blacklist:jti:{jti}"
                return bool(self.redis.exists(key))
            except Exception as e:
                print(f"[TokenManager] Redis error during is_jti_blacklisted: {e}. Using in-memory fallback.")

        # Fallback check
        now = time.time()
        with self._mem_lock:
            exp = self._memory_blacklist.get(jti)
            if exp and exp > now:
                return True
            elif exp and exp <= now:
                del self._memory_blacklist[jti]
        return False

    def get_jti_status(self, jti: str) -> Dict[str, Any]:
        """Returns details on whether a JTI is blacklisted and its remaining TTL."""
        if not self.use_fallback and self.redis:
            try:
                key = f"blacklist:jti:{jti}"
                ttl = self.redis.ttl(key)
                is_blacklisted = ttl > 0 or self.redis.exists(key) == 1
                return {
                    "jti": jti,
                    "is_blacklisted": is_blacklisted,
                    "ttl_seconds": max(0, ttl) if is_blacklisted else 0,
                    "storage": "redis"
                }
            except Exception:
                pass

        now = time.time()
        with self._mem_lock:
            exp = self._memory_blacklist.get(jti)
            if exp and exp > now:
                return {
                    "jti": jti,
                    "is_blacklisted": True,
                    "ttl_seconds": int(exp - now),
                    "storage": "in_memory"
                }
        return {
            "jti": jti,
            "is_blacklisted": False,
            "ttl_seconds": 0,
            "storage": "in_memory"
        }

    # ---------------------------------------------------------------------------
    # Session Tracking & Mass Revocation
    # ---------------------------------------------------------------------------
    def track_user_jti(self, sub: str, jti: str, exp_timestamp: float):
        """Registers active JTI under user's session list for bulk logout capability."""
        if not self.use_fallback and self.redis:
            try:
                user_key = f"user:{sub}:jtis"
                val = json.dumps({"jti": jti, "exp": exp_timestamp})
                self.redis.sadd(user_key, val)
                # Keep user session set alive for 7 days
                self.redis.expire(user_key, REFRESH_TOKEN_EXPIRE_DAYS * 86400)
                return
            except Exception:
                pass

        with self._mem_lock:
            if sub not in self._memory_user_jtis:
                self._memory_user_jtis[sub] = []
            self._memory_user_jtis[sub].append((jti, exp_timestamp))

    def revoke_all_user_jtis(self, sub: str) -> int:
        """Mass revokes all active JTIs for a specified user (Logout-All)."""
        revoked_count = 0
        now = time.time()

        if not self.use_fallback and self.redis:
            try:
                user_key = f"user:{sub}:jtis"
                members = self.redis.smembers(user_key)
                for member in members:
                    try:
                        data = json.loads(member)
                        jti = data.get("jti")
                        exp = data.get("exp", now + 3600)
                        if jti and exp > now:
                            self.revoke_jti(jti, exp)
                            revoked_count += 1
                    except Exception:
                        continue
                self.redis.delete(user_key)
                return revoked_count
            except Exception:
                pass

        with self._mem_lock:
            jtis = self._memory_user_jtis.pop(sub, [])
            for jti, exp in jtis:
                if exp > now:
                    self._memory_blacklist[jti] = exp
                    revoked_count += 1
        return revoked_count

    def _purge_expired_memory_blacklist(self):
        """Internal helper to prune expired JTIs from memory dictionary."""
        now = time.time()
        expired = [jti for jti, exp in self._memory_blacklist.items() if exp <= now]
        for jti in expired:
            del self._memory_blacklist[jti]