app/auth/config.py

from fastapi_users.authentication import (
    AuthenticationBackend,
    BearerTransport,
    CookieTransport,
    JWTStrategy,
)
from app.config import get_settings
from typing import Optional, Any
import jwt
from jwt import PyJWK
from datetime import datetime, timedelta

settings = get_settings()

# Bearer transport (para APIs programáticas / mobile)
bearer_transport = BearerTransport(tokenUrl="api/auth/jwt/login")

# Cookie transport (para navegadores - HttpOnly, Secure, SameSite=Lax)
cookie_transport = CookieTransport(
    cookie_name="cd_token",
    cookie_max_age=settings.access_token_expire_minutes * 60,
    cookie_secure=not settings.debug,  # Secure=True en prod, False en dev local
    cookie_httponly=True,
    cookie_samesite="lax",
)

# ─── Custom Dual JWT Strategy ───
# Firma nuevos tokens con RS256 (private key)
# Verifica: intenta RS256 (public key) → si falla, fallback a HS256 (legacy secret)

class DualJWTStrategy(JWTStrategy):
    """JWT Strategy que soporta verificación dual: RS256 (nuevo) + HS256 (legacy)."""

    def __init__(self):
        # Cargar claves
        self._private_key = self._load_private_key()
        self._public_key = self._load_public_key()
        self._legacy_secret = settings.jwt_legacy_secret_key or settings.secret_key
        self._algorithm = settings.jwt_algorithm
        self._legacy_algorithm = settings.jwt_legacy_algorithm
        self._lifetime_seconds = settings.access_token_expire_minutes * 60
        self._key_id = settings.jwt_key_id
        
        # Initialize base class with required params (using legacy secret for base compat)
        super().__init__(
            secret=self._legacy_secret,
            lifetime_seconds=self._lifetime_seconds,
            token_audience=["fastapi-users:auth"],
            algorithm=self._algorithm,
            public_key=self._public_key,
        )
    
    def _load_private_key(self) -> Optional[str]:
        """Cargar clave privada RS256 desde config o archivo."""
        if settings.jwt_private_key:
            return settings.jwt_private_key
        # Fallback: leer archivo
        import os
        key_path = os.path.join(os.path.dirname(__file__), '..', '..', 'private_key.pem')
        if os.path.exists(key_path):
            with open(key_path, 'r') as f:
                return f.read()
        return None
    
    def _load_public_key(self) -> Optional[str]:
        """Cargar clave pública RS256 desde config o archivo."""
        if settings.jwt_public_key:
            return settings.jwt_public_key
        # Fallback: leer archivo
        import os
        key_path = os.path.join(os.path.dirname(__file__), '..', '..', 'public_key.pem')
        if os.path.exists(key_path):
            with open(key_path, 'r') as f:
                return f.read()
        return None

    async def write_token(self, user) -> str:
        """Firmar token con RS256 (nueva clave privada). Fallar si no hay clave."""
        # Extraer ID del usuario (fastapi-users pasa el objeto User)
        user_id = str(user.id)
        data = {"sub": user_id, "aud": self.token_audience}
        
        if not self._private_key:
            raise RuntimeError(
                "JWT_PRIVATE_KEY no configurada. Configure RS256 keys para firmar tokens. "
                "No se permite fallback silencioso a HS256."
            )
        
        # Agregar kid en header para rotación de claves
        headers = {"kid": self._key_id}
        return jwt.encode(
            {**data, "exp": datetime.utcnow() + timedelta(seconds=self._lifetime_seconds)},
            self._private_key,
            algorithm=self._algorithm,
            headers=headers,
        )

    async def read_token(self, token: Optional[str], user_manager) -> Optional[Any]:
        """Verificar token: solo RS256. Legacy HS256 solo si jwt_legacy_enabled=True."""
        if token is None:
            return None
        
        # Intentar RS256 (nuevo)
        if self._public_key:
            try:
                data = jwt.decode(
                    token,
                    self._public_key,
                    algorithms=[self._algorithm],
                    audience=self.token_audience,
                )
                user_id = data.get("sub")
                if user_id is None:
                    return None
                parsed_id = user_manager.parse_id(user_id)
                return await user_manager.get(parsed_id)
            except jwt.PyJWTError:
                pass  # Fallar a legacy si habilitado
        
        # Legacy HS256 SOLO si explícitamente habilitado
        if settings.jwt_legacy_enabled and self._legacy_secret:
            try:
                data = jwt.decode(
                    token,
                    self._legacy_secret,
                    algorithms=[self._legacy_algorithm],
                    audience=self.token_audience,
                )
                user_id = data.get("sub")
                if user_id is None:
                    return None
                parsed_id = user_manager.parse_id(user_id)
                return await user_manager.get(parsed_id)
            except jwt.PyJWTError:
                pass
        
        # Si llegamos aquí, token inválido
        return None

    async def destroy_token(self, token: str, user) -> None:
        """JWT es stateless, no hay nada que destruir server-side."""
        pass


def get_jwt_strategy() -> DualJWTStrategy:
    return DualJWTStrategy()


# Dual auth backend: soporta AMBOS transports (header Authorization + cookie)
auth_backend = AuthenticationBackend(
    name="jwt",
    transport=bearer_transport,  # primary para compatibilidad
    get_strategy=get_jwt_strategy,
)

# Segundo backend solo para cookies (se usa en login para setear cookie)
cookie_auth_backend = AuthenticationBackend(
    name="jwt-cookie",
    transport=cookie_transport,
    get_strategy=get_jwt_strategy,
)