"""TOTP (Time-based One-Time Password) 2FA endpoints."""
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
import pyotp
import auth

router = APIRouter()

class Setup2FAResponse(BaseModel):
    secret: str
    uri: str

class Verify2FARequest(BaseModel):
    secret: str
    code: str

@router.post("/setup-2fa/generate", response_model=Setup2FAResponse)
async def generate_2fa(current_user = Depends(auth.get_current_user)):
    """Generate a new TOTP secret and provisioning URI for 2FA setup."""
    secret = pyotp.random_base32()
    uri = pyotp.totp.TOTP(secret).provisioning_uri(
        name=current_user.username,
        issuer_name="NAS Dashboard"
    )
    return {"secret": secret, "uri": uri}

@router.post("/setup-2fa/verify")
async def verify_2fa_setup(request: Verify2FARequest, current_user = Depends(auth.get_current_user)):
    """Verify TOTP code and enable 2FA for the user."""
    totp = pyotp.TOTP(request.secret)
    if not totp.verify(request.code):
        raise HTTPException(status_code=400, detail="Invalid code")

    auth.save_totp_secret(request.secret)
    return {"message": "2FA enabled successfully"}

@router.post("/setup-2fa/disable")
async def disable_2fa(current_user = Depends(auth.get_current_user)):
    """Disable 2FA for the user."""
    auth.save_totp_secret("")
    return {"message": "2FA disabled"}