rmi-backend/app/domains/auth/__init__.py

"""
Auth Router - Complete authentication system (email, wallet, OAuth, Telegram)
"""

import json
import logging
from datetime import datetime, timedelta
from typing import Any, cast

from fastapi import APIRouter, HTTPException, Request

from app.core.redis import get_redis
from app.domains.auth.deps import (
    get_current_user,
    require_auth as require_auth,
    require_public_profile as require_public_profile,
)
from app.domains.auth.jwt import (
    JWT_ALGORITHM,
    JWT_EXPIRY_DAYS as JWT_EXPIRY_DAYS,
    JWT_SECRET,
    _create_jwt,
    _derive_user_id,
    _verify_jwt,
    generate_nonce,
)
from app.domains.auth.oauth import router as oauth_router
from app.domains.auth.passwords import hash_password, verify_password
from app.domains.auth.schemas import (
    EmailLoginRequest,
    EmailRegisterRequest,
    GoogleAuthResponse as GoogleAuthResponse,
    NonceResponse,
    TelegramAuthRequest as TelegramAuthRequest,
    TwoFAEnableRequest,
    TwoFALoginRequest,
    TwoFASetupResponse as TwoFASetupResponse,
    TwoFAStatusResponse,
    TwoFAVerifyRequest,
    UserResponse,
    WalletAuthResponse,
    WalletNonceRequest,
    WalletVerifyRequest,
)
from app.domains.auth.store import (
    _delete_user as _delete_user,
    _get_user,
    _get_user_by_email,
    _is_valid_email,
    _save_user,
)
from app.domains.auth.totp import (
    PYOTP_AVAILABLE,
    TOTP_DIGITS as TOTP_DIGITS,
    TOTP_INTERVAL as TOTP_INTERVAL,
    TOTP_ISSUER as TOTP_ISSUER,
    TOTP_WINDOW as TOTP_WINDOW,
    _build_totp as _build_totp,
    _generate_backup_codes,
    _generate_qr_base64,
    _generate_totp_secret,
    _get_totp_uri,
    _hash_backup_code,
    _verify_backup_code,
    _verify_totp,
)

logger = logging.getLogger(__name__)

router = APIRouter(tags=["auth"])

# Include OAuth/Telegram routes from dedicated submodule
router.include_router(oauth_router, tags=["auth"])

# ── Config ──


# ── Storage (Redis-backed user store) ──
# ── Email Auth ──
@router.post("/register", response_model=WalletAuthResponse)
def register_email(req: EmailRegisterRequest) -> WalletAuthResponse:
    """Register a new user with email/password."""
    if not _is_valid_email(req.email):
        raise HTTPException(status_code=400, detail="Invalid email format")

    if len(req.password) < 8:
        raise HTTPException(status_code=400, detail="Password must be at least 8 characters")

    # Check if user exists
    existing = _get_user_by_email(req.email)
    if existing:
        raise HTTPException(status_code=400, detail="Email already registered")

    user_id = _derive_user_id(req.email)
    hashed = hash_password(req.password)

    user: dict[str, Any] = {
        "id": user_id,
        "email": req.email,
        "display_name": req.display_name,
        "password_hash": hashed,
        "wallet_address": req.wallet_address,
        "wallet_chain": req.wallet_chain,
        "tier": "FREE",
        "role": "USER",
        "created_at": datetime.utcnow().isoformat(),
        "xp": 0,
        "level": 1,
        "badges": [],
        "scans_remaining": 5,
        "scans_used": 0,
    }

    # Save user + email index
    r = get_redis()
    assert r is not None
    r.hset("rmi:users", user_id, json.dumps(user))
    r.hset("rmi:users:email", req.email.lower(), user_id)

    token = _create_jwt(user_id, req.email, "FREE", "USER", req.wallet_address)

    return cast(
        "WalletAuthResponse",
        {
            "access_token": token,
            "refresh_token": token,
            "user": {
                "id": user_id,
                "email": req.email,
                "display_name": req.display_name,
                "wallet_address": req.wallet_address,
                "wallet_chain": req.wallet_chain,
                "tier": "FREE",
                "role": "USER",
                "created_at": user["created_at"],
            },
        },
    )


@router.post("/login", response_model=WalletAuthResponse)
def login_email(req: EmailLoginRequest) -> WalletAuthResponse:
    """Login with email/password. If 2FA enabled, returns partial token requiring /2fa/login."""
    user = _get_user_by_email(req.email)
    if not user or not user.get("password_hash"):
        raise HTTPException(status_code=401, detail="Invalid credentials")

    if not verify_password(req.password, user["password_hash"]):
        raise HTTPException(status_code=401, detail="Invalid credentials")

    # If 2FA enabled, return a pre-auth token that requires TOTP verification
    if user.get("totp_enabled") and user.get("totp_secret"):
        # Generate a short-lived pre-auth token (5 min)
        pre_token = _create_jwt(
            user["id"],
            user["email"],
            user.get("tier", "FREE"),
            user.get("role", "USER"),
            user.get("wallet_address"),
        )
        # Override expiry to 5 minutes
        import jose.jwt as jose_jwt

        payload = jose_jwt.decode(pre_token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        payload["exp"] = datetime.utcnow() + timedelta(minutes=5)
        payload["2fa_required"] = True
        payload["pre_auth"] = True
        pre_token = jose_jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)

        return cast(
            "WalletAuthResponse",
            {
                "access_token": pre_token,
                "refresh_token": pre_token,
                "user": {
                    "id": user["id"],
                    "email": user["email"],
                    "display_name": user.get("display_name", user["email"]),
                    "wallet_address": user.get("wallet_address"),
                    "wallet_chain": user.get("wallet_chain"),
                    "tier": user.get("tier", "FREE"),
                    "role": user.get("role", "USER"),
                    "created_at": user.get("created_at"),
                    "_2fa_required": True,
                },
            },
        )

    token = _create_jwt(
        user["id"],
        user["email"],
        user.get("tier", "FREE"),
        user.get("role", "USER"),
        user.get("wallet_address"),
    )

    return cast(
        "WalletAuthResponse",
        {
            "access_token": token,
            "refresh_token": token,
            "user": {
                "id": user["id"],
                "email": user["email"],
                "display_name": user.get("display_name", user["email"]),
                "wallet_address": user.get("wallet_address"),
                "wallet_chain": user.get("wallet_chain"),
                "tier": user.get("tier", "FREE"),
                "role": user.get("role", "USER"),
                "created_at": user.get("created_at"),
            },
        },
    )


# ── Wallet Auth ──
@router.post("/wallet/nonce", response_model=NonceResponse)
async def wallet_nonce(req: WalletNonceRequest) -> NonceResponse:
    """Get a nonce for wallet signature."""
    nonce = generate_nonce()
    timestamp = datetime.utcnow().isoformat()
    message = f"RugMunch Intelligence wants you to sign in with your {req.chain.title()} account.\n\nWallet: {req.address}\nNonce: {nonce}\nTimestamp: {timestamp}"
    return cast("NonceResponse", {"nonce": nonce, "timestamp": timestamp, "message": message})


@router.post("/wallet/verify", response_model=WalletAuthResponse)
async def wallet_verify(req: WalletVerifyRequest) -> WalletAuthResponse:
    """Verify wallet signature and create session."""
    from app.domains.auth.wallet import get_or_create_wallet_user, verify_wallet_signature

    if not verify_wallet_signature(req.message, req.signature, req.address):
        raise HTTPException(status_code=401, detail="Invalid signature")

    user_data = await get_or_create_wallet_user(req.address)

    return cast(
        "WalletAuthResponse",
        {
            "access_token": user_data["access_token"],
            "refresh_token": user_data["refresh_token"],
            "user": {
                "id": user_data["id"],
                "email": user_data["email"],
                "display_name": user_data.get("display_name", f"Agent {req.address[2:8].upper()}"),
                "wallet_address": req.address,
                "wallet_chain": req.chain,
                "tier": user_data["tier"],
                "role": user_data["role"],
                "created_at": user_data["created_at"],
            },
        },
    )


# ── User Profile ──
@router.get("/user/me", response_model=UserResponse)
async def user_me(request: Request) -> UserResponse:
    """Get current authenticated user profile."""
    auth_header = request.headers.get("Authorization", "")
    if not auth_header.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing auth token")

    token = auth_header[7:]
    payload = _verify_jwt(token)
    if not payload:
        raise HTTPException(status_code=401, detail="Invalid token")

    user = _get_user(payload["id"])
    if not user:
        raise HTTPException(status_code=404, detail="User not found")

    return cast(
        "UserResponse",
        {
            "id": user["id"],
            "email": user["email"],
            "display_name": user.get("display_name", user["email"]),
            "wallet_address": user.get("wallet_address"),
            "wallet_chain": user.get("wallet_chain"),
            "tier": user.get("tier", "FREE"),
            "role": user.get("role", "USER"),
            "created_at": user.get("created_at"),
            "xp": user.get("xp", 0),
            "level": user.get("level", 1),
            "badges": user.get("badges", []),
        },
    )


# ── 2FA / TOTP Endpoints ──


@router.post("/2fa/setup")
async def twofa_setup(request: Request) -> TwoFASetupResponse:
    """Begin 2FA setup - generate secret + QR code. Returns plaintext secret + backup codes (shown once)."""
    if not PYOTP_AVAILABLE:
        raise HTTPException(status_code=503, detail="2FA service unavailable")

    user = await get_current_user(request)
    if not user:
        raise HTTPException(status_code=401, detail="Authentication required")

    # Check if already enabled
    if user.get("totp_secret"):
        raise HTTPException(
            status_code=400, detail="2FA already enabled. Disable first to reconfigure."
        )

    secret = _generate_totp_secret()
    uri = _get_totp_uri(secret, user["email"])
    qr_b64 = _generate_qr_base64(uri)
    backup_codes = _generate_backup_codes(8)

    # Store pending secret (not enabled until verified)
    r = get_redis()
    assert r is not None
    pending = {
        "secret": secret,
        "backup_codes": [_hash_backup_code(c) for c in backup_codes],
        "created_at": datetime.utcnow().isoformat(),
    }
    r.setex(f"rmi:2fa_pending:{user['id']}", timedelta(minutes=10), json.dumps(pending))

    return cast("TwoFASetupResponse", {
        "secret": secret,
        "qr_code": qr_b64,
        "uri": uri,
        "backup_codes": backup_codes,  # plaintext - shown once
    })


@router.post("/2fa/enable")
async def twofa_enable(req: TwoFAEnableRequest, request: Request) -> dict[str, Any]:
    """Enable 2FA - verify the first TOTP code, then persist."""
    if not PYOTP_AVAILABLE:
        raise HTTPException(status_code=503, detail="2FA service unavailable")

    user = await get_current_user(request)
    if not user:
        raise HTTPException(status_code=401, detail="Authentication required")

    r = get_redis()
    assert r is not None
    pending_raw = r.get(f"rmi:2fa_pending:{user['id']}")
    if not pending_raw:
        raise HTTPException(status_code=400, detail="No pending 2FA setup. Call /2fa/setup first.")

    pending = json.loads(pending_raw)
    secret = pending["secret"]

    # Verify the code
    if not _verify_totp(secret, req.code):
        raise HTTPException(
            status_code=400, detail="Invalid TOTP code. Check your authenticator app time sync."
        )

    # Enable 2FA on user record
    user["totp_secret"] = secret
    user["totp_enabled"] = True
    user["totp_created_at"] = pending["created_at"]
    user["totp_backup_codes"] = pending["backup_codes"]  # already hashed
    user["totp_last_used"] = None
    _save_user(user)

    # Clean up pending
    r.delete(f"rmi:2fa_pending:{user['id']}")

    # Audit log
    logger.info(f"[2FA ENABLED] user={user['id']} email={user['email']}")

    return {"status": "ok", "message": "2FA enabled successfully", "method": "totp"}


@router.post("/2fa/disable")
async def twofa_disable(request: Request) -> dict[str, Any]:
    """Disable 2FA - requires current password for security."""
    user = await get_current_user(request)
    if not user:
        raise HTTPException(status_code=401, detail="Authentication required")

    if not user.get("totp_enabled"):
        raise HTTPException(status_code=400, detail="2FA is not enabled")

    # Remove 2FA fields
    for key in [
        "totp_secret",
        "totp_enabled",
        "totp_created_at",
        "totp_backup_codes",
        "totp_last_used",
    ]:
        user.pop(key, None)

    _save_user(user)
    logger.info(f"[2FA DISABLED] user={user['id']} email={user['email']}")

    return cast("dict[str, Any]", {"status": "ok", "message": "2FA disabled successfully"})


@router.get("/2fa/status", response_model=TwoFAStatusResponse)
async def twofa_status(request: Request) -> TwoFAStatusResponse:
    """Check 2FA status for current user."""
    user = await get_current_user(request)
    if not user:
        raise HTTPException(status_code=401, detail="Authentication required")

    enabled = user.get("totp_enabled", False)
    return cast("TwoFAStatusResponse", {
        "enabled": enabled,
        "method": "totp" if enabled else "none",
        "created_at": user.get("totp_created_at"),
        "last_used": user.get("totp_last_used"),
    })


@router.post("/2fa/verify")
async def twofa_verify(req: TwoFAVerifyRequest, request: Request) -> dict[str, Any]:
    """Verify a TOTP code (for re-auth during sensitive operations)."""
    if not PYOTP_AVAILABLE:
        raise HTTPException(status_code=503, detail="2FA service unavailable")

    user = await get_current_user(request)
    if not user:
        raise HTTPException(status_code=401, detail="Authentication required")

    if not user.get("totp_enabled") or not user.get("totp_secret"):
        raise HTTPException(status_code=400, detail="2FA not enabled")

    if not _verify_totp(user["totp_secret"], req.code):
        raise HTTPException(status_code=400, detail="Invalid TOTP code")

    # Update last used
    user["totp_last_used"] = datetime.utcnow().isoformat()
    _save_user(user)

    return cast("dict[str, Any]", {"status": "ok", "verified": True})


@router.post("/2fa/login")
async def twofa_login(req: TwoFALoginRequest) -> WalletAuthResponse:
    """Login with email + password + 2FA code. Returns full JWT on success."""
    if not PYOTP_AVAILABLE:
        raise HTTPException(status_code=503, detail="2FA service unavailable")

    user = _get_user_by_email(req.email)
    if not user or not user.get("password_hash"):
        raise HTTPException(status_code=401, detail="Invalid credentials")

    if not verify_password(req.password, user["password_hash"]):
        raise HTTPException(status_code=401, detail="Invalid credentials")

    # Check if 2FA is enabled
    if not user.get("totp_enabled") or not user.get("totp_secret"):
        raise HTTPException(
            status_code=400, detail="2FA not enabled for this account. Use /login instead."
        )

    code = req.code.strip().replace("-", "")

    # Try TOTP first
    if _verify_totp(user["totp_secret"], code):
        pass  # valid TOTP
    else:
        # Try backup codes
        backup_codes = user.get("totp_backup_codes", [])
        matched = False
        for i, hashed in enumerate(backup_codes):
            if _verify_backup_code(req.code, hashed):
                matched = True
                # Remove used backup code
                backup_codes.pop(i)
                user["totp_backup_codes"] = backup_codes
                break
        if not matched:
            raise HTTPException(status_code=401, detail="Invalid 2FA code or backup code")

    # Update last used
    user["totp_last_used"] = datetime.utcnow().isoformat()
    _save_user(user)

    token = _create_jwt(
        user["id"],
        user["email"],
        user.get("tier", "FREE"),
        user.get("role", "USER"),
        user.get("wallet_address"),
    )

    logger.info(f"[2FA LOGIN] user={user['id']} email={user['email']}")

    return cast(
        "WalletAuthResponse",
        {
            "access_token": token,
            "refresh_token": token,
            "user": {
                "id": user["id"],
                "email": user["email"],
                "display_name": user.get("display_name", user["email"]),
                "wallet_address": user.get("wallet_address"),
                "wallet_chain": user.get("wallet_chain"),
                "tier": user.get("tier", "FREE"),
                "role": user.get("role", "USER"),
                "created_at": user.get("created_at"),
            },
        },
    )