arrr-erp/app/ui/routes/auth.py

from __future__ import annotations

from datetime import datetime, timezone, timedelta

from fastapi import APIRouter, Form, Request
from fastapi.responses import RedirectResponse
from sqlalchemy import select

from app.core.db.common import CommonSessionLocal
from app.core.security.csrf import get_or_create_csrf_token, validate_csrf
from app.core.security.otp import start_otp, verify_otp
from app.core.security.passwords import verify_password, hash_password
from app.core.security.session_auth import (
    SESSION_LOGIN_AT_KEY,
    SESSION_USER_ID_KEY,
    get_current_user,
)
from app.core.templating import templates
from app.core.settings import get_settings
from app.modules.core.iam.invite_service import accept_invite, reset_password_with_token
from app.modules.core.iam.models import LoginAttempt, User
from app.modules.core.rbac.models import Permission, Role, RolePermission, UserRole
from app.modules.core.tenancy.models import Branch, FinancialYear, Tenant
from app.modules.core.tenancy.settings_models import BranchSettings
from app.modules.email_integration.services import send_auth_otp_email, send_password_changed_email

router = APIRouter()

PENDING_POST_LOGIN_REDIRECT_KEY = "pending_post_login_redirect"
SAFE_POST_LOGIN_REDIRECTS = {
    "/mobile/attendance",
    "/employee/attendance",
    "/employee/dashboard",
}


def _consume_safe_post_login_redirect(request: Request) -> str | None:
    value = request.session.pop(PENDING_POST_LOGIN_REDIRECT_KEY, None)
    value = (value or "").strip()
    if value in SAFE_POST_LOGIN_REDIRECTS:
        return value
    return None


def _dev_otp_print_enabled() -> bool:
    settings = get_settings()
    return bool(getattr(settings, "DEV_AUTH_OTP_PRINT", False)) and (settings.ENV or "").lower() in {"dev", "local", "development"}


def _log_dev_otp(label: str, email: str, code: str) -> None:
    if _dev_otp_print_enabled():
        print(f"[DEV OTP] {label} user={email} code={code}")


def _client_ip(request: Request) -> str:
    return request.client.host if request.client else "unknown"


def _attempt_key(email: str, ip: str) -> str:
    return f"{email.lower().strip()}|{ip}"


def _get_branch_security_policy(db, user: User) -> BranchSettings | None:
    if not getattr(user, "branch_id", None):
        return None
    return db.execute(
        select(BranchSettings).where(BranchSettings.branch_id == user.branch_id)
    ).scalar_one_or_none()


def _user_roles(db, user_id: int) -> list[str]:
    q = (
        select(Role.name)
        .join(UserRole, UserRole.role_id == Role.id)
        .where(UserRole.user_id == user_id)
    )
    return [r for (r,) in db.execute(q).all()]


def _user_permissions(db, user_id: int) -> set[str]:
    q = (
        select(Permission.code)
        .join(RolePermission, RolePermission.permission_id == Permission.id)
        .join(UserRole, UserRole.role_id == RolePermission.role_id)
        .where(UserRole.user_id == user_id, Permission.is_active.is_(True))
    )
    return set(db.execute(q).scalars().all())


def _otp_required(bs: BranchSettings | None, roles: list[str]) -> bool:
    if not bs:
        return False
    required = {
        x.strip() for x in (bs.otp_required_roles_csv or "").split(",") if x.strip()
    }
    return any(r in required for r in roles)

def _default_financial_year_code(db, tenant_id: int | None) -> str | None:
    if not tenant_id:
        return None
    fy = db.execute(
        select(FinancialYear).where(
            FinancialYear.tenant_id == tenant_id,
            FinancialYear.is_current.is_(True),
        )
    ).scalar_one_or_none()
    if fy:
        return fy.year_code
    fy = db.execute(
        select(FinancialYear)
        .where(FinancialYear.tenant_id == tenant_id)
        .order_by(FinancialYear.start_date.desc(), FinancialYear.year_code.desc())
    ).scalars().first()
    return fy.year_code if fy else None


def _tenant_code(db, tenant_id: int | None) -> str | None:
    if not tenant_id:
        return None
    tenant = db.get(Tenant, int(tenant_id))
    return tenant.code if tenant else None


def _branch_code(db, branch_id: int | None) -> str | None:
    if not branch_id:
        return None
    branch = db.get(Branch, int(branch_id))
    return branch.code if branch else None


def _post_login_redirect(must_change_password: bool, permissions: set[str], roles: list[str]) -> str:
    if must_change_password:
        return "/change-password-required"

    role_set = set(roles or [])

    # Platform owner always lands on platform control centre first.
    if "System Admin" in role_set:
        return "/system-admin/dashboard"

    if "Client" in role_set:
        return "/client/dashboard"

    if "Consultant" in role_set:
        return "/consultant/dashboard"

    # If Firm Admin is also Partner, daily operations are more frequent;
    # the workspace switcher exposes Firm Administration when required.
    if "Firm Admin" in role_set and "Partner" in role_set:
        return "/partner/dashboard"

    if "Firm Admin" in role_set:
        return "/firm-admin/dashboard"

    if "Partner" in role_set:
        return "/partner/dashboard"

    if role_set.intersection({"Manager", "Branch Manager"}):
        return "/manager/dashboard"

    if (
        "employees.ess.view" in permissions
        or "employees.work.view_self" in permissions
        or "employees.attendance.view_self" in permissions
        or "employees.leave.view_self" in permissions
        or "employees.documents.view_self" in permissions
        or "employees.payroll.view_self" in permissions
        or "Staff" in role_set
    ):
        return "/employee/dashboard"

    if "system.settings.view" in permissions or "users.view" in permissions:
        return "/system-settings"

    return "/employee/dashboard"

def _is_user_login_allowed(user: User) -> tuple[bool, str | None]:
    if not user:
        return False, "Invalid credentials"

    if not getattr(user, "is_active", True):
        return False, "User account is inactive."

    if hasattr(user, "allow_login") and not bool(getattr(user, "allow_login", True)):
        return False, "Login is disabled for this account."

    if hasattr(user, "is_locked") and bool(getattr(user, "is_locked", False)):
        return False, "User account is locked."

    if hasattr(user, "deleted_at") and getattr(user, "deleted_at", None) is not None:
        return False, "User account is deleted."

    return True, None


def _template_context(
    request: Request,
    db=None,
    *,
    title: str,
    flash: str | None = None,
    extra: dict | None = None,
) -> dict:
    ctx = {
        "request": request,
        "csrf_token": get_or_create_csrf_token(request),
        "flash": flash,
        "title": title,
    }

    if db is not None:
        current_user = get_current_user(request, db=db)
        if current_user:
            ctx.update(
                {
                    "current_user": current_user,
                    "current_user_roles": _user_roles(db, int(current_user.id)),
                    "current_user_permissions": list(
                        _user_permissions(db, int(current_user.id))
                    ),
                }
            )

    if extra:
        ctx.update(extra)

    return ctx


def _render_login(request: Request, flash: str | None = None, status_code: int = 200):
    return templates.TemplateResponse(
        "modules/core/iam/templates/login.html",
        _template_context(request, title="Login", flash=flash),
        status_code=status_code,
    )


def _render_otp(request: Request, flash: str | None = None, status_code: int = 200):
    db = CommonSessionLocal()
    try:
        return templates.TemplateResponse(
            "modules/core/iam/templates/otp.html",
            _template_context(
                request,
                db=db,
                title="OTP Verification",
                flash=flash,
            ),
            status_code=status_code,
        )
    finally:
        db.close()


def _render_change_password(
    request: Request, flash: str | None = None, status_code: int = 200
):
    db = CommonSessionLocal()
    try:
        return templates.TemplateResponse(
            "modules/core/iam/templates/change_password.html",
            _template_context(
                request,
                db=db,
                title="Change Password",
                flash=flash,
            ),
            status_code=status_code,
        )
    finally:
        db.close()


def _render_change_password_otp(
    request: Request, flash: str | None = None, status_code: int = 200
):
    db = CommonSessionLocal()
    try:
        return templates.TemplateResponse(
            "modules/core/iam/templates/change_password_otp.html",
            _template_context(
                request,
                db=db,
                title="Confirm Password Change",
                flash=flash,
            ),
            status_code=status_code,
        )
    finally:
        db.close()


def _render_forgot_password(
    request: Request, flash: str | None = None, status_code: int = 200
):
    return templates.TemplateResponse(
        "modules/core/iam/templates/forgot_password.html",
        _template_context(
            request,
            title="Forgot Password",
            flash=flash,
        ),
        status_code=status_code,
    )


def _render_reset_password(
    request: Request, flash: str | None = None, status_code: int = 200
):
    return templates.TemplateResponse(
        "modules/core/iam/templates/reset_password.html",
        _template_context(
            request,
            title="Reset Password",
            flash=flash,
        ),
        status_code=status_code,
    )


def _render_invite_accept(
    request: Request, token: str, flash: str | None = None, status_code: int = 200
):
    return templates.TemplateResponse(
        "modules/core/iam/templates/invite_accept.html",
        _template_context(
            request,
            title="Accept Invite",
            flash=flash,
            extra={"token": token},
        ),
        status_code=status_code,
    )


@router.get("/invite/accept")
def invite_accept_page(request: Request, token: str = ""):
    if not token.strip():
        return RedirectResponse(url="/login", status_code=303)
    return _render_invite_accept(request, token=token.strip())


@router.post("/invite/accept")
def invite_accept_submit(
    request: Request,
    token: str = Form(...),
    password: str = Form(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    token_clean = token.strip()
    if not token_clean:
        return RedirectResponse(url="/login", status_code=303)
    db = CommonSessionLocal()
    try:
        try:
            user = accept_invite(db, token_clean, password.strip())
        except ValueError as exc:
            return _render_invite_accept(request, token=token_clean, flash=str(exc), status_code=400)
        if not user:
            return _render_invite_accept(request, token=token_clean, flash="Invalid or expired invite link.", status_code=400)
        return RedirectResponse(url="/login", status_code=303)
    finally:
        db.close()


@router.get("/login")
def login_page(request: Request):
    return _render_login(request)


@router.post("/login")
def login_submit(
    request: Request,
    email: str = Form(...),
    password: str = Form(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)

    email_clean = email.strip().lower()
    ip = _client_ip(request)
    key = _attempt_key(email_clean, ip)

    db = CommonSessionLocal()
    try:
        la = db.execute(
            select(LoginAttempt).where(LoginAttempt.key == key)
        ).scalar_one_or_none()
        now = datetime.now(timezone.utc)

        if la and la.locked_until_utc and la.locked_until_utc.replace(
            tzinfo=timezone.utc
        ) > now:
            return _render_login(
                request,
                flash=f"Account temporarily locked. Try again after {la.locked_until_utc}.",
                status_code=429,
            )

        user = db.execute(
            select(User).where(User.email == email_clean)
        ).scalar_one_or_none()

        can_login, blocked_reason = _is_user_login_allowed(user)
        password_ok = bool(user and verify_password(password, user.password_hash))

        if not user or not can_login or not password_ok:
            lock_attempts = 5
            lock_minutes = 15

            if user:
                bs = _get_branch_security_policy(db, user)
                if bs:
                    lock_attempts = bs.lockout_attempts
                    lock_minutes = bs.lockout_minutes

            if not la:
                la = LoginAttempt(key=key, attempts=0, updated_at_utc=now)
                db.add(la)

            la.attempts = int(la.attempts or 0) + 1
            la.updated_at_utc = now

            if la.attempts >= lock_attempts:
                la.locked_until_utc = now + timedelta(minutes=lock_minutes)
                la.attempts = 0

            db.commit()

            flash = blocked_reason or "Invalid credentials"
            return _render_login(request, flash=flash, status_code=400)

        if la:
            la.attempts = 0
            la.locked_until_utc = None
            la.updated_at_utc = now
            db.commit()

        user_id = int(user.id)
        user_email = str(user.email)
        tenant_id = getattr(user, "tenant_id", None)
        branch_id = getattr(user, "branch_id", None)
        must_change_password = bool(getattr(user, "must_change_password", False))

        roles = _user_roles(db, user_id)
        permissions = _user_permissions(db, user_id)
        bs = _get_branch_security_policy(db, user)

        request.session[SESSION_USER_ID_KEY] = user_id
        request.session[SESSION_LOGIN_AT_KEY] = now.isoformat()
        request.session["user_email"] = user_email
        tenant_code = _tenant_code(db, tenant_id)
        branch_code = _branch_code(db, branch_id)

        request.session["tenant_id"] = tenant_id
        request.session["branch_id"] = branch_id
        request.session["active_tenant_id"] = tenant_id
        request.session["active_branch_id"] = branch_id
        if tenant_code:
            request.session["tenant_code"] = tenant_code
            request.session["active_tenant_code"] = tenant_code
        else:
            request.session.pop("tenant_code", None)
            request.session.pop("active_tenant_code", None)
        if branch_code:
            request.session["branch_code"] = branch_code
            request.session["active_branch_code"] = branch_code
        else:
            request.session.pop("branch_code", None)
            request.session.pop("active_branch_code", None)

        active_financial_year = _default_financial_year_code(db, tenant_id)
        if active_financial_year:
            request.session["active_financial_year"] = active_financial_year
        request.session["must_change_password"] = must_change_password
        default_post_login_redirect = _post_login_redirect(
            must_change_password, permissions, roles
        )
        pending_post_login_redirect = None if must_change_password else _consume_safe_post_login_redirect(request)
        request.session["post_login_redirect"] = pending_post_login_redirect or default_post_login_redirect

        if _otp_required(bs, roles):
            code = start_otp(request)
            try:
                send_auth_otp_email(db, user=user, otp_code=code, purpose="login")
                db.commit()
            except Exception as exc:
                print(f"[EMAIL OTP ERROR] user={user_email} error={exc}")
            _log_dev_otp("login", user_email, code)
            request.session["otp_verified"] = False
            return RedirectResponse(url="/otp", status_code=303)

        request.session["otp_verified"] = True
        return RedirectResponse(
            url=request.session.get("post_login_redirect", "/system-settings"),
            status_code=303,
        )

    finally:
        db.close()


@router.get("/otp")
def otp_page(request: Request):
    if not request.session.get(SESSION_USER_ID_KEY):
        return RedirectResponse(url="/login", status_code=303)
    return _render_otp(request)


@router.post("/otp")
def otp_submit(
    request: Request,
    otp: str = Form(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)

    if not request.session.get(SESSION_USER_ID_KEY):
        return RedirectResponse(url="/login", status_code=303)

    if verify_otp(request, otp):
        request.session["otp_verified"] = True
        return RedirectResponse(
            url=request.session.get("post_login_redirect", "/system-settings"),
            status_code=303,
        )

    return _render_otp(
        request,
        flash="Invalid OTP. Please check the OTP sent to your registered email.",
        status_code=400,
    )


@router.get("/change-password")
def change_password_page(request: Request):
    if not request.session.get(SESSION_USER_ID_KEY):
        return RedirectResponse(url="/login", status_code=303)
    return _render_change_password(request)


@router.post("/change-password")
def change_password_submit(
    request: Request,
    current_password: str = Form(...),
    new_password: str = Form(...),
    confirm_password: str = Form(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)

    user_id = request.session.get(SESSION_USER_ID_KEY)
    if not user_id:
        return RedirectResponse(url="/login", status_code=303)

    if new_password != confirm_password:
        return _render_change_password(
            request,
            flash="New password and confirm password do not match.",
            status_code=400,
        )

    if len(new_password.strip()) < 8:
        return _render_change_password(
            request,
            flash="New password must be at least 8 characters.",
            status_code=400,
        )

    db = CommonSessionLocal()
    try:
        user = db.execute(select(User).where(User.id == user_id)).scalar_one_or_none()
        if not user:
            request.session.clear()
            return RedirectResponse(url="/login", status_code=303)

        if not verify_password(current_password, user.password_hash):
            return _render_change_password(
                request,
                flash="Current password is incorrect.",
                status_code=400,
            )

        request.session["pending_password_change_hash"] = hash_password(
            new_password.strip()
        )
        request.session["pending_password_change_user_id"] = int(user.id)

        code = start_otp(request)
        try:
            send_auth_otp_email(db, user=user, otp_code=code, purpose="password_change")
            db.commit()
        except Exception as exc:
            print(f"[EMAIL OTP ERROR] password-change user={user.email} error={exc}")
        _log_dev_otp("password-change", str(user.email), code)

        return RedirectResponse(url="/change-password/otp", status_code=303)
    finally:
        db.close()


@router.get("/change-password/otp")
def change_password_otp_page(request: Request):
    if not request.session.get(SESSION_USER_ID_KEY):
        return RedirectResponse(url="/login", status_code=303)

    if not request.session.get("pending_password_change_hash"):
        return RedirectResponse(url="/change-password", status_code=303)

    return _render_change_password_otp(request)


@router.post("/change-password/otp")
def change_password_otp_submit(
    request: Request,
    otp: str = Form(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)

    user_id = request.session.get(SESSION_USER_ID_KEY)
    pending_user_id = request.session.get("pending_password_change_user_id")
    pending_hash = request.session.get("pending_password_change_hash")

    if not user_id:
        return RedirectResponse(url="/login", status_code=303)

    if not pending_hash or not pending_user_id or int(user_id) != int(pending_user_id):
        return RedirectResponse(url="/change-password", status_code=303)

    if not verify_otp(request, otp):
        return _render_change_password_otp(
            request,
            flash="Invalid OTP. Please check the OTP sent to your registered email.",
            status_code=400,
        )

    db = CommonSessionLocal()
    try:
        user = db.execute(select(User).where(User.id == user_id)).scalar_one_or_none()
        if not user:
            request.session.clear()
            return RedirectResponse(url="/login", status_code=303)

        user.password_hash = pending_hash
        user.must_change_password = False
        user.password_changed_at_utc = datetime.now(timezone.utc)
        try:
            send_password_changed_email(db, user=user)
        except Exception as exc:
            print(f"[EMAIL PASSWORD CHANGED ERROR] user={user.email} error={exc}")
        db.commit()

        roles = _user_roles(db, int(user.id))
        permissions = _user_permissions(db, int(user.id))

        request.session.pop("pending_password_change_hash", None)
        request.session.pop("pending_password_change_user_id", None)
        request.session["must_change_password"] = False
        request.session["post_login_redirect"] = _post_login_redirect(False, permissions, roles)

        return RedirectResponse(url=request.session.get("post_login_redirect", "/system-settings"), status_code=303)
    finally:
        db.close()


@router.get("/forgot-password")
def forgot_password_page(request: Request):
    return _render_forgot_password(request)


@router.post("/forgot-password")
def forgot_password_submit(
    request: Request,
    email: str = Form(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)

    email_clean = email.strip().lower()

    db = CommonSessionLocal()
    try:
        user = db.execute(
            select(User).where(User.email == email_clean)
        ).scalar_one_or_none()

        request.session.pop("password_reset_user_id", None)
        request.session.pop("password_reset_email", None)

        if not user:
            return _render_forgot_password(
                request,
                flash="If the login ID exists, password reset instructions have been sent to the registered email.",
                status_code=200,
            )

        can_login, _ = _is_user_login_allowed(user)
        if not can_login:
            return _render_forgot_password(
                request,
                flash="If the login ID exists, password reset instructions have been sent to the registered email.",
                status_code=200,
            )

        request.session["password_reset_user_id"] = int(user.id)
        request.session["password_reset_email"] = str(user.email)

        code = start_otp(request)
        try:
            send_auth_otp_email(db, user=user, otp_code=code, purpose="password_reset")
            db.commit()
        except Exception as exc:
            print(f"[EMAIL OTP ERROR] forgot-password user={user.email} error={exc}")
        _log_dev_otp("forgot-password", str(user.email), code)

        return RedirectResponse(url="/reset-password", status_code=303)
    finally:
        db.close()


@router.get("/reset-password")
def reset_password_page(request: Request):
    if not request.session.get("password_reset_user_id"):
        return RedirectResponse(url="/forgot-password", status_code=303)
    return _render_reset_password(request)


@router.post("/reset-password")
def reset_password_submit(
    request: Request,
    otp: str = Form(...),
    new_password: str = Form(...),
    confirm_password: str = Form(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)

    user_id = request.session.get("password_reset_user_id")
    if not user_id:
        return RedirectResponse(url="/forgot-password", status_code=303)

    if not verify_otp(request, otp):
        return _render_reset_password(
            request,
            flash="Invalid OTP. Please check the OTP sent to your registered email.",
            status_code=400,
        )

    if new_password != confirm_password:
        return _render_reset_password(
            request,
            flash="New password and confirm password do not match.",
            status_code=400,
        )

    if len(new_password.strip()) < 8:
        return _render_reset_password(
            request,
            flash="New password must be at least 8 characters.",
            status_code=400,
        )

    db = CommonSessionLocal()
    try:
        user = db.execute(select(User).where(User.id == user_id)).scalar_one_or_none()
        if not user:
            request.session.pop("password_reset_user_id", None)
            request.session.pop("password_reset_email", None)
            return RedirectResponse(url="/forgot-password", status_code=303)

        user.password_hash = hash_password(new_password.strip())
        user.must_change_password = False
        user.password_changed_at_utc = datetime.now(timezone.utc)
        try:
            send_password_changed_email(db, user=user)
        except Exception as exc:
            print(f"[EMAIL PASSWORD CHANGED ERROR] user={user.email} error={exc}")
        db.commit()

        request.session.pop("password_reset_user_id", None)
        request.session.pop("password_reset_email", None)

        return RedirectResponse(url="/login", status_code=303)
    finally:
        db.close()


@router.get("/password-reset/accept")
def password_reset_token_page(request: Request, token: str = ""):
    return templates.TemplateResponse(
        "modules/core/iam/templates/reset_password_token.html",
        _template_context(request, title="Reset Password", extra={"token": token.strip()}),
    )


@router.post("/password-reset/accept")
def password_reset_token_submit(
    request: Request,
    token: str = Form(...),
    new_password: str = Form(...),
    confirm_password: str = Form(...),
    csrf_token: str = Form(...),
):
    validate_csrf(request, csrf_token)
    token_clean = token.strip()
    if new_password != confirm_password:
        return templates.TemplateResponse(
            "modules/core/iam/templates/reset_password_token.html",
            _template_context(request, title="Reset Password", flash="New password and confirm password do not match.", extra={"token": token_clean}),
            status_code=400,
        )
    if len(new_password.strip()) < get_settings().PASSWORD_MIN_LENGTH:
        return templates.TemplateResponse(
            "modules/core/iam/templates/reset_password_token.html",
            _template_context(request, title="Reset Password", flash=f"New password must be at least {get_settings().PASSWORD_MIN_LENGTH} characters.", extra={"token": token_clean}),
            status_code=400,
        )
    db = CommonSessionLocal()
    try:
        try:
            user = reset_password_with_token(db, token_clean, new_password.strip())
        except ValueError as exc:
            return templates.TemplateResponse(
                "modules/core/iam/templates/reset_password_token.html",
                _template_context(request, title="Reset Password", flash=str(exc), extra={"token": token_clean}),
                status_code=400,
            )
        if not user:
            return templates.TemplateResponse(
                "modules/core/iam/templates/reset_password_token.html",
                _template_context(request, title="Reset Password", flash="Invalid or expired password reset link.", extra={"token": token_clean}),
                status_code=400,
            )
        try:
            send_password_changed_email(db, user=user)
        except Exception as exc:
            print(f"[EMAIL PASSWORD CHANGED ERROR] user={getattr(user, 'email', '')} error={exc}")
        db.commit()
        return RedirectResponse(url="/login", status_code=303)
    finally:
        db.close()


@router.get("/change-password-required")
def change_password_required(request: Request):
    return RedirectResponse(url="/change-password", status_code=303)


@router.get("/logout")
def logout(request: Request):
    request.session.clear()
    return RedirectResponse(url="/login", status_code=303)