admin / Strike

from datetime import datetime, timezone

from fastapi import APIRouter, Depends, HTTPException, status, Request
from fastapi.security import OAuth2PasswordRequestForm
from jose import jwt, JWTError
from sqlalchemy.orm import Session

from .. import models, schemas
from ..config import settings
from ..database import get_db
from ..security import (
    verify_password, create_access_token, client_ip, write_audit,
    get_current_user, verify_totp, hash_password, password_problem,
)

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


def _token_response(user: models.User) -> schemas.Token:
    token = create_access_token(user.username, scope="access")
    return schemas.Token(
        access_token=token,
        username=user.username,
        role=user.role,
        must_change_password=user.must_change_password,
        mfa_required=user.mfa_required,
        mfa_enabled=user.mfa_enabled,
        first_login_pending=user.first_login_pending,
    )


@router.post("/signup", status_code=201)
def signup(payload: schemas.SignupCreate, request: Request,
           db: Session = Depends(get_db)):
    """
    Public self-service signup. Creates a new account with the ``user`` role.

    The new user can only ever see the assessments they create. MFA is enforced
    on their first sign-in (``mfa_required`` + ``first_login_pending``).
    """
    ip = client_ip(request)
    surname = (payload.surname or "").strip()
    first_name = (payload.first_name or "").strip()
    company = (payload.company_name or "").strip()
    email = (payload.email or "").strip()
    contact = (payload.contact_number or "").strip()
    username = (payload.username or "").strip()

    missing = [n for n, v in [
        ("Surname", surname), ("First Name", first_name),
        ("Company Name", company), ("Email Address", email),
        ("Contact Number", contact), ("Username", username),
    ] if not v]
    if missing:
        raise HTTPException(status_code=400,
                            detail=f"Required field(s) missing: {', '.join(missing)}.")
    if "@" not in email or "." not in email.split("@")[-1]:
        raise HTTPException(status_code=400, detail="Please enter a valid email address.")

    prob = password_problem(payload.password)
    if prob:
        raise HTTPException(status_code=400, detail=prob)

    if db.query(models.User).filter(models.User.username == username).first():
        raise HTTPException(status_code=409, detail=f"Username '{username}' is already taken.")

    user = models.User(
        username=username,
        password_hash=hash_password(payload.password),
        full_name=f"{first_name} {surname}".strip(),
        role="user",
        is_active=True,
        must_change_password=False,
        first_name=first_name,
        surname=surname,
        company_name=company,
        email=email,
        contact_number=contact,
        signup_origin="self_service",
        # Enforce MFA enrolment on first login.
        mfa_required=True,
        mfa_enabled=False,
        first_login_pending=True,
    )
    db.add(user)
    db.commit()
    write_audit(db, username, "SIGNUP",
                f"Self-service signup company='{company}' email='{email}'", ip)
    return {"detail": "Account created. Please sign in; you'll be asked to set up "
                      "multi-factor authentication on first login."}


@router.post("/login")
def login(request: Request, form: OAuth2PasswordRequestForm = Depends(),
          db: Session = Depends(get_db)):
    """
    Step 1 of login: validate username + password.

    * If the user has MFA enabled, returns a short-lived ``mfa_token`` and
      ``mfa_required: true`` instead of an access token; the client must then
      call ``/login/mfa`` with a TOTP code.
    * Otherwise returns a normal access token.
    """
    ip = client_ip(request)
    user = db.query(models.User).filter(models.User.username == form.username).first()
    if not user or not verify_password(form.password, user.password_hash):
        write_audit(db, form.username, "LOGIN_FAILED", "Invalid credentials", ip)
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED,
                            detail="Invalid username or password")
    if not user.is_active:
        write_audit(db, form.username, "LOGIN_FAILED", "Account disabled", ip)
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Account disabled")

    if user.mfa_enabled and user.mfa_secret:
        # Password OK but a second factor is required.
        mfa_token = create_access_token(user.username, scope="mfa", minutes=5)
        write_audit(db, user.username, "LOGIN_MFA_CHALLENGE",
                    "Password OK; awaiting TOTP", ip)
        return schemas.MfaChallenge(mfa_token=mfa_token, username=user.username)

    user.last_login = datetime.now(timezone.utc)
    # A self-service account on its very first login has MFA required but not yet
    # enabled; the client routes them straight to MFA enrolment. We leave
    # first_login_pending set until MFA is actually confirmed (see users router).
    db.commit()
    if user.mfa_required and not user.mfa_enabled:
        write_audit(db, user.username, "LOGIN_SUCCESS",
                    "Password only; MFA enrolment required", ip)
    else:
        write_audit(db, user.username, "LOGIN_SUCCESS", "Password only", ip)
    return _token_response(user)


@router.post("/login/mfa", response_model=schemas.Token)
def login_mfa(request: Request, payload: dict, db: Session = Depends(get_db)):
    """Step 2 of login: verify the TOTP code against the mfa_token subject."""
    ip = client_ip(request)
    mfa_token = payload.get("mfa_token", "")
    code = payload.get("code", "")
    cred_exc = HTTPException(status_code=status.HTTP_401_UNAUTHORIZED,
                             detail="MFA verification failed")
    try:
        data = jwt.decode(mfa_token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
        if data.get("scope") != "mfa":
            raise cred_exc
        username = data.get("sub")
    except JWTError:
        raise cred_exc

    user = db.query(models.User).filter(models.User.username == username).first()
    if not user or not user.is_active or not user.mfa_enabled:
        raise cred_exc
    if not verify_totp(user.mfa_secret, code):
        write_audit(db, username, "LOGIN_MFA_FAILED", "Bad TOTP code", ip)
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED,
                            detail="Incorrect authentication code")

    user.last_login = datetime.now(timezone.utc)
    if user.first_login_pending:
        user.first_login_pending = False
    db.commit()
    write_audit(db, user.username, "LOGIN_SUCCESS", "Password + MFA", ip)
    return _token_response(user)


@router.post("/logout")
def logout(request: Request, db: Session = Depends(get_db),
           user: models.User = Depends(get_current_user)):
    write_audit(db, user.username, "LOGOUT", "", client_ip(request))
    return {"detail": "Logged out"}


@router.get("/me")
def me(user: models.User = Depends(get_current_user)):
    return {
        "username": user.username,
        "role": user.role,
        "full_name": user.full_name,
        "must_change_password": user.must_change_password,
        "mfa_enabled": user.mfa_enabled,
        "mfa_required": user.mfa_required,
        "first_login_pending": user.first_login_pending,
    }