app/auth/router.py

import uuid
import json
import secrets
from typing import Optional
from datetime import datetime
from fastapi import APIRouter, Depends, HTTPException, Response, Request, Cookie
from fastapi_users import FastAPIUsers, exceptions as fu_exceptions
from fastapi_users.jwt import generate_jwt
from app.auth.models import User, LoginHistory
from app.auth.manager import get_user_manager, UserManager
from app.auth.config import auth_backend, cookie_auth_backend
from app.auth.schemas import UserRead, UserCreate, UserUpdate
from app.database import get_db
from sqlalchemy.ext.asyncio import AsyncSession
from pydantic import BaseModel
import traceback
import pyotp
import qrcode
import io
import base64
from app.config import get_settings
from app.utils.security import mask_email
from app.auth.models import User as UserModel

settings = get_settings()

# Incluir AMBOS backends para dual auth (header + cookie)
fastapi_users = FastAPIUsers[User, uuid.UUID](
    get_user_manager,
    [auth_backend, cookie_auth_backend],
)

current_active_user = fastapi_users.current_user(active=True)

router = APIRouter(prefix="/auth", tags=["auth"])

# Auth router con AMBOS backends
router.include_router(
    fastapi_users.get_auth_router(auth_backend, requires_verification=False),
    prefix="/jwt",
)

router.include_router(
    fastapi_users.get_auth_router(cookie_auth_backend, requires_verification=False),
    prefix="/jwt",
)

router.include_router(
    fastapi_users.get_register_router(UserRead, UserCreate),
)

router.include_router(
    fastapi_users.get_users_router(UserRead, UserUpdate),
    prefix="/users",
)


class ForgotPasswordRequest(BaseModel):
    email: str


class ResetPasswordRequest(BaseModel):
    token: str
    password: str


@router.post("/forgot-password")
async def forgot_password(
    body: ForgotPasswordRequest,
    manager: UserManager = Depends(get_user_manager),
):
    try:
        user = await manager.get_by_email(body.email)
    except fu_exceptions.UserNotExists:
        return {"message": "Si el email está registrado, se envió un enlace de recuperación."}

    token_data = {
        "sub": str(user.id),
        "password_fgpt": manager.password_helper.hash(user.hashed_password),
        "aud": manager.reset_password_token_audience,
    }
    token = generate_jwt(
        token_data,
        manager.reset_password_token_secret,
        manager.reset_password_token_lifetime_seconds,
    )
    await manager.on_after_forgot_password(user, token, None)

    return {
        "message": "Si el email está registrado, se envió un enlace de recuperación.",
    }


@router.post("/reset-password")
async def reset_password(
    body: ResetPasswordRequest,
    manager: UserManager = Depends(get_user_manager),
):
    try:
        user = await manager.reset_password(body.token, body.password)
    except fu_exceptions.InvalidResetPasswordToken:
        raise HTTPException(status_code=400, detail="Token inválido o expirado.")
    except fu_exceptions.UserInactive:
        raise HTTPException(status_code=400, detail="Usuario inactivo.")

    await manager.on_after_reset_password(user, None)
    return {"message": "Contraseña actualizada correctamente."}


# Custom login que setea cookie HttpOnly Y devuelve token en body (compatibilidad)
@router.post("/jwt/login")
async def jwt_login(
    request: Request,
    response: Response,
    manager: UserManager = Depends(get_user_manager),
):
    try:
        # Parse form data (OAuth2 password flow)
        form = await request.form()
        username = form.get("username")
        password = form.get("password")
        
        if not username or not password:
            raise HTTPException(status_code=400, detail="username y password requeridos")
        
        # Autenticar usuario usando OAuth2PasswordRequestForm
        from fastapi.security import OAuth2PasswordRequestForm
        credentials = OAuth2PasswordRequestForm(username=username, password=password)
        user = await manager.authenticate(credentials)
        if not user:
            raise HTTPException(status_code=400, detail="Credenciales inválidas")
        if not user.is_active:
            raise HTTPException(status_code=400, detail="Usuario inactivo")
        
        # Si MFA está habilitado, requerir desafío TOTP/backup code
        if user.mfa_enabled:
            # Generar token temporal de pre-autenticación (válido 5 min)
            from fastapi_users.jwt import generate_jwt
            from app.config import get_settings
            settings = get_settings()
            mfa_token_data = {
                "sub": str(user.id),
                "mfa_pending": True,
                "aud": "mfa-challenge",
            }
            mfa_token = generate_jwt(
                mfa_token_data,
                settings.secret_key,
                300,  # 5 minutos
            )
            return {
                "mfa_required": True,
                "mfa_token": mfa_token,
                "message": "Introduce tu código TOTP o código de respaldo",
            }
        
        # Generar access token
        jwt_strategy = auth_backend.get_strategy()
        access_token = await jwt_strategy.write_token(user)
        
        # Crear refresh token (rotación)
        from app.database import get_db
        async for db in get_db():
            refresh_token = await manager.create_refresh_token(user, request, db)
        
        # Setear cookie HttpOnly para access token
        from app.config import get_settings
        settings = get_settings()
        cookie_max_age = settings.access_token_expire_minutes * 60
        response.set_cookie(
            key="cd_token",
            value=access_token,
            max_age=cookie_max_age,
            httponly=True,
            secure=not settings.debug,
            samesite="lax",
            path="/",
        )
        
        # Setear cookie HttpOnly para refresh token
        refresh_max_age = settings.refresh_token_expire_days * 24 * 60 * 60
        response.set_cookie(
            key="cd_refresh_token",
            value=refresh_token,
            max_age=refresh_max_age,
            httponly=True,
            secure=not settings.debug,
            samesite="lax",
            path="/",
        )
        
        # También setear cookie de expiración para que el frontend pueda leerla
        import time
        response.set_cookie(
            key="cd_token_expiry",
            value=str(int(time.time() * 1000) + cookie_max_age * 1000),
            max_age=cookie_max_age,
            httponly=True,
            secure=not settings.debug,
            samesite="lax",
            path="/",
        )
        
        # Devolver token en body para compatibilidad con clientes existentes
        return {
            "access_token": access_token,
            "token_type": "bearer",
        }
    except HTTPException:
        raise
    except Exception as e:
        import logging
        logger = logging.getLogger("crowdata.auth")
        logger.error(f"Error en jwt_login: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Error interno: {str(e)}")


# Refresh token endpoint - rota access token + nuevo refresh token
@router.post("/jwt/refresh")
async def jwt_refresh(
    request: Request,
    response: Response,
    cd_refresh_token: Optional[str] = Cookie(None, alias="cd_refresh_token"),
    manager: UserManager = Depends(get_user_manager),
):
    if not cd_refresh_token:
        raise HTTPException(status_code=401, detail="Refresh token requerido")
    
    # Verificar refresh token
    from app.database import get_db
    async for db in get_db():
        user = await manager.verify_refresh_token(cd_refresh_token, db)
    
    if not user:
        raise HTTPException(status_code=401, detail="Refresh token inválido o expirado")
    
    # Generar nuevo access token
    jwt_strategy = auth_backend.get_strategy()
    access_token = await jwt_strategy.write_token(user)
    
    # Rotar refresh token (revocar viejo, crear nuevo)
    from app.database import get_db
    async for db in get_db():
        new_refresh_token = await manager.create_refresh_token(user, request, db)
    
    # Setear cookies
    from app.config import get_settings
    settings = get_settings()
    cookie_max_age = settings.access_token_expire_minutes * 60
    refresh_max_age = settings.refresh_token_expire_days * 24 * 60 * 60
    
    response.set_cookie(
        key="cd_token",
        value=access_token,
        max_age=cookie_max_age,
        httponly=True,
        secure=not settings.debug,
        samesite="lax",
        path="/",
    )
    response.set_cookie(
        key="cd_refresh_token",
        value=new_refresh_token,
        max_age=refresh_max_age,
        httponly=True,
        secure=not settings.debug,
        samesite="lax",
        path="/",
    )
    import time
    response.set_cookie(
        key="cd_token_expiry",
        value=str(int(time.time() * 1000) + cookie_max_age * 1000),
        max_age=cookie_max_age,
        httponly=True,
        secure=not settings.debug,
        samesite="lax",
        path="/",
    )
    
    return {
        "access_token": access_token,
        "token_type": "bearer",
    }


# Logout: limpiar cookies
@router.post("/jwt/logout")
async def jwt_logout(
    response: Response,
    cd_refresh_token: Optional[str] = Cookie(None, alias="cd_refresh_token"),
    manager: UserManager = Depends(get_user_manager),
):
    # Revocar refresh token en BD
    if cd_refresh_token:
        from app.database import get_db
        async for db in get_db():
            await manager.revoke_refresh_token(cd_refresh_token, db)
    
    response.delete_cookie("cd_token", path="/", httponly=True, secure=True, samesite="lax")
    response.delete_cookie("cd_refresh_token", path="/", httponly=True, secure=True, samesite="lax")
    response.delete_cookie("cd_token_expiry", path="/", httponly=True, secure=True, samesite="lax")
    return {"message": "Sesión cerrada"}


# ─── JWKS Endpoint para claves públicas (RFC 7517) ───
@router.get("/.well-known/jwks.json", tags=["auth"])
async def jwks():
    """JSON Web Key Set - expone claves públicas para verificación RS256."""
    from app.config import get_settings
    settings = get_settings()
    
    if not 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:
                public_key_pem = f.read()
        else:
            raise HTTPException(status_code=503, detail="Clave pública no configurada")
    else:
        public_key_pem = settings.jwt_public_key
    
    # Convertir PEM a JWK
    from jwt.algorithms import RSAAlgorithm
    public_key = RSAAlgorithm.from_jwk(public_key_pem)
    numbers = public_key.public_numbers()
    
    # Base64url encode sin padding
    import base64
    def b64url_encode(data: bytes) -> str:
        return base64.urlsafe_b64encode(data).decode().rstrip('=')
    
    n = b64url_encode(numbers.n.to_bytes((numbers.n.bit_length() + 7) // 8, 'big'))
    e = b64url_encode(numbers.e.to_bytes((numbers.e.bit_length() + 7) // 8, 'big'))
    
    return {
        "keys": [
            {
                "kty": "RSA",
                "use": "sig",
                "alg": "RS256",
                "kid": settings.jwt_key_id,
                "n": n,
                "e": e,
            }
        ]
    }


# ════════════════════════════════════════════════════════════════════════════
# MFA / 2FA (TOTP) Endpoints
# ════════════════════════════════════════════════════════════════════════════

class MFASetupRequest(BaseModel):
    password: str  # Confirmar contraseña actual


class MFAVerifyRequest(BaseModel):
    code: str  # Código TOTP de 6 dígitos


class MFADisableRequest(BaseModel):
    password: str
    code: str


@router.post("/mfa/setup", dependencies=[Depends(current_active_user)])
async def mfa_setup(
    response: Response,
    request: Request,
    body: MFASetupRequest,
    user: User = Depends(current_active_user),
    manager: UserManager = Depends(get_user_manager),
):
    """
    Iniciar configuración de MFA.
    Genera secreto TOTP y devuelve QR code (base64) + secret para app autenticadora.
    Requiere contraseña actual para confirmar identidad.
    """
    # Verificar contraseña
    if not manager.password_helper.verify(body.password, user.hashed_password):
        raise HTTPException(status_code=400, detail="Contraseña incorrecta")
    
    if user.mfa_enabled:
        raise HTTPException(status_code=400, detail="MFA ya está habilitado")
    
    import pyotp
    import qrcode
    import io
    import base64
    
    # Generar secreto único
    secret = pyotp.random_base32()
    
    # Generar URI para QR code (compatible con Google Authenticator, Authy, etc.)
    totp_uri = pyotp.totp.TOTP(secret).provisioning_uri(
        name=user.email,
        issuer_name="CrowData",
    )
    
    # Generar QR code como base64
    qr = qrcode.QRCode(version=1, box_size=10, border=5)
    qr.add_data(totp_uri)
    qr.make(fit=True)
    img = qr.make_image(fill_color="black", back_color="white")
    
    buf = io.BytesIO()
    img.save(buf, format='PNG')
    qr_base64 = base64.b64encode(buf.getvalue()).decode()
    
    # Guardar secreto temporal (no activar hasta verificar) - using encrypted property
    user.mfa_secret_decrypted = secret
    user.mfa_backup_codes_decrypted = []
    from app.database import get_db
    async for db in get_db():
        await db.commit()
    
    return {
        "secret": secret,
        "qr_code": f"data:image/png;base64,{qr_base64}",
        "uri": totp_uri,
        "message": "Escanea el QR con tu app autenticadora (Google Authenticator, Authy, 1Password, etc.) y luego usa /mfa/verify para activar."
    }


@router.post("/mfa/verify", dependencies=[Depends(current_active_user)])
async def mfa_verify(
    body: MFAVerifyRequest,
    user: User = Depends(current_active_user),
    manager: UserManager = Depends(get_user_manager),
):
    """
    Verificar código TOTP y activar MFA.
    Genera códigos de respaldo (backup codes) al activar.
    """
    if not user.mfa_secret_decrypted:
        raise HTTPException(status_code=400, detail="MFA no iniciado. Usa /mfa/setup primero.")
    
    if user.mfa_enabled:
        raise HTTPException(status_code=400, detail="MFA ya está habilitado")
    
    import pyotp
    
    totp = pyotp.TOTP(user.mfa_secret_decrypted)
    if not totp.verify(body.code, valid_window=1):
        raise HTTPException(status_code=400, detail="Código inválido o expirado")
    
    # Generar backup codes (8 códigos de 8 chars cada uno)
    import secrets
    backup_codes = [secrets.token_urlsafe(6) for _ in range(8)]
    
    user.mfa_enabled = True
    user.mfa_backup_codes_decrypted = backup_codes
    user.mfa_verified_at = datetime.utcnow()
    
    from app.database import get_db
    async for db in get_db():
        await db.commit()
    
    return {
        "message": "MFA activado correctamente",
        "backup_codes": backup_codes,
        "warning": "Guarda estos códigos de respaldo en un lugar seguro. Cada uno se puede usar una sola vez."
    }


@router.post("/mfa/disable", dependencies=[Depends(current_active_user)])
async def mfa_disable(
    body: MFADisableRequest,
    user: User = Depends(current_active_user),
    manager: UserManager = Depends(get_user_manager),
):
    """
    Desactivar MFA.
    Requiere contraseña + código TOTP actual (o backup code).
    """
    if not user.mfa_enabled:
        raise HTTPException(status_code=400, detail="MFA no está habilitado")
    
    # Verificar contraseña
    if not manager.password_helper.verify(body.password, user.hashed_password):
        raise HTTPException(status_code=400, detail="Contraseña incorrecta")
    
    # Verificar código (TOTP o backup code)
    import pyotp
    
    valid = False
    totp = pyotp.TOTP(user.mfa_secret_decrypted)
    
    if totp.verify(body.code, valid_window=1):
        valid = True
    elif body.code in user.mfa_backup_codes_decrypted:
        # Es un backup code - consumirlo
        codes = user.mfa_backup_codes_decrypted
        codes.remove(body.code)
        user.mfa_backup_codes_decrypted = codes
        valid = True
    
    if not valid:
        raise HTTPException(status_code=400, detail="Código inválido")
    
    # Desactivar MFA
    user.mfa_enabled = False
    user.mfa_secret_decrypted = None
    user.mfa_backup_codes_decrypted = []
    user.mfa_verified_at = None
    
    from app.database import get_db
    async for db in get_db():
        await db.commit()
    
    return {"message": "MFA desactivado correctamente"}


@router.post("/mfa/backup-codes", dependencies=[Depends(current_active_user)])
async def mfa_regenerate_backup_codes(
    user: User = Depends(current_active_user),
    manager: UserManager = Depends(get_user_manager),
):
    """
    Regenerar códigos de respaldo (invalida los anteriores).
    """
    if not user.mfa_enabled:
        raise HTTPException(status_code=400, detail="MFA no está habilitado")
    
    import secrets
    backup_codes = [secrets.token_urlsafe(6) for _ in range(8)]
    
    user.mfa_backup_codes_decrypted = backup_codes
    from app.database import get_db
    async for db in get_db():
        await db.commit()
    
    return {
        "backup_codes": backup_codes,
        "warning": "Los códigos anteriores han sido invalidados. Guarda los nuevos en un lugar seguro."
    }


@router.get("/mfa/status", dependencies=[Depends(current_active_user)])
async def mfa_status(user: User = Depends(current_active_user)):
    """Estado actual de MFA del usuario."""
    return {
        "mfa_enabled": user.mfa_enabled,
        "mfa_verified_at": user.mfa_verified_at,
        "backup_codes_remaining": len(user.mfa_backup_codes_decrypted),
    }


class MFAChallengeRequest(BaseModel):
    mfa_token: str  # Token temporal del login inicial
    code: str       # Código TOTP o backup code


@router.post("/mfa/challenge")
async def mfa_challenge(
    body: MFAChallengeRequest,
    response: Response,
    manager: UserManager = Depends(get_user_manager),
):
    """
    Verificar código TOTP/backup code durante login con MFA habilitado.
    Recibe el token temporal (mfa_token) del login inicial + código.
    Si es válido, setea cookies y devuelve access_token.
    """
    # Verificar token temporal MFA
    from app.config import get_settings
    settings = get_settings()
    try:
        import jwt as pyjwt
        payload = pyjwt.decode(
            body.mfa_token,
            settings.secret_key,
            algorithms=["HS256"],
            audience="mfa-challenge",
        )
    except pyjwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Token MFA inválido o expirado")
    
    if not payload.get("mfa_pending"):
        raise HTTPException(status_code=401, detail="Token MFA inválido")
    
    user_id = payload.get("sub")
    if not user_id:
        raise HTTPException(status_code=401, detail="Token MFA inválido")
    
    import uuid
    user = await manager.get(uuid.UUID(user_id))
    if not user or not user.is_active:
        raise HTTPException(status_code=401, detail="Usuario no encontrado o inactivo")
    
    if not user.mfa_enabled:
        raise HTTPException(status_code=400, detail="MFA no está habilitado para este usuario")
    
    # Verificar código TOTP
    import pyotp
    totp = pyotp.TOTP(user.mfa_secret_decrypted)
    
    code_valid = False
    if totp.verify(body.code, valid_window=1):
        code_valid = True
    elif body.code in user.mfa_backup_codes_decrypted:
        # Es un backup code - consumirlo
        codes = user.mfa_backup_codes_decrypted
        codes.remove(body.code)
        user.mfa_backup_codes_decrypted = codes
        code_valid = True
    
    if not code_valid:
        raise HTTPException(status_code=400, detail="Código inválido o expirado")
    
    # Código válido - generar tokens normales
    jwt_strategy = auth_backend.get_strategy()
    access_token = await jwt_strategy.write_token(user)
    
    # Crear refresh token
    from app.database import get_db
    from starlette.requests import Request
    async for db in get_db():
        refresh_token = await manager.create_refresh_token(user, Request({"type": "http"}), db)
    
    # Setear cookies HttpOnly
    cookie_max_age = settings.access_token_expire_minutes * 60
    refresh_max_age = settings.refresh_token_expire_days * 24 * 60 * 60
    
    response.set_cookie(
        key="cd_token",
        value=access_token,
        max_age=cookie_max_age,
        httponly=True,
        secure=not settings.debug,
        samesite="lax",
        path="/",
    )
    response.set_cookie(
        key="cd_refresh_token",
        value=refresh_token,
        max_age=refresh_max_age,
        httponly=True,
        secure=not settings.debug,
        samesite="lax",
        path="/",
    )
    import time
    response.set_cookie(
        key="cd_token_expiry",
        value=str(int(time.time() * 1000) + cookie_max_age * 1000),
        max_age=cookie_max_age,
        httponly=True,
        secure=not settings.debug,
        samesite="lax",
        path="/",
    )
    
    return {
        "access_token": access_token,
        "token_type": "bearer",
        "message": "Autenticación MFA completada"
    }