resolutionflow/backend/app/api/endpoints/onboarding.py

"""Onboarding status endpoints — track new-user checklist progress."""

from typing import Annotated

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.deps import get_current_active_user
from app.core.database import get_db
from app.core.admin_database import get_admin_db
from app.models.account import Account
from app.models.assistant_chat import AssistantChat
from app.models.psa_connection import PsaConnection
from app.models.session import Session
from app.models.tree import Tree
from app.models.user import User
from app.schemas.onboarding import (
    OnboardingStatus,
    OnboardingStepRequest,
    OnboardingStepResponse,
)

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


@router.get("/onboarding-status", response_model=OnboardingStatus)
async def get_onboarding_status(
    db: Annotated[AsyncSession, Depends(get_db)],
    current_user: Annotated[User, Depends(get_current_active_user)],
) -> OnboardingStatus:
    """Return onboarding checklist completion status for the current user."""

    # created_flow — user owns at least 1 tree
    created_flow_q = await db.execute(
        select(func.count()).select_from(Tree).where(Tree.author_id == current_user.id).limit(1)
    )
    created_flow = (created_flow_q.scalar() or 0) > 0

    # ran_session — user has at least 1 session
    ran_session_q = await db.execute(
        select(func.count()).select_from(Session).where(Session.user_id == current_user.id).limit(1)
    )
    ran_session = (ran_session_q.scalar() or 0) > 0

    # exported_session — user has at least 1 session with exported=True
    exported_q = await db.execute(
        select(func.count())
        .select_from(Session)
        .where(Session.user_id == current_user.id, Session.exported == True)  # noqa: E712
        .limit(1)
    )
    exported_session = (exported_q.scalar() or 0) > 0

    # tried_ai_assistant — user has at least 1 assistant chat
    ai_q = await db.execute(
        select(func.count())
        .select_from(AssistantChat)
        .where(AssistantChat.user_id == current_user.id)
        .limit(1)
    )
    tried_ai_assistant = (ai_q.scalar() or 0) > 0

    # Team-dependent checks
    is_team_user = current_user.team_id is not None

    invited_teammate = False
    connected_psa = False

    if is_team_user:
        # invited_teammate — team/account has more than 1 member
        if current_user.account_id:
            teammate_q = await db.execute(
                select(func.count())
                .select_from(User)
                .where(
                    User.account_id == current_user.account_id,
                    User.deleted_at.is_(None),
                )
            )
            invited_teammate = (teammate_q.scalar() or 0) > 1

        # connected_psa — account has at least 1 PSA connection
        if current_user.account_id:
            psa_q = await db.execute(
                select(func.count())
                .select_from(PsaConnection)
                .where(PsaConnection.account_id == current_user.account_id)
                .limit(1)
            )
            connected_psa = (psa_q.scalar() or 0) > 0

    return OnboardingStatus(
        created_flow=created_flow,
        ran_session=ran_session,
        exported_session=exported_session,
        tried_ai_assistant=tried_ai_assistant,
        invited_teammate=invited_teammate,
        connected_psa=connected_psa,
        is_team_user=is_team_user,
        dismissed=current_user.onboarding_dismissed,
    )


@router.post("/onboarding-status/dismiss", response_model=OnboardingStatus)
async def dismiss_onboarding(
    db: Annotated[AsyncSession, Depends(get_admin_db)],
    current_user: Annotated[User, Depends(get_current_active_user)],
) -> OnboardingStatus:
    """Dismiss the onboarding checklist for the current user."""
    current_user.onboarding_dismissed = True
    await db.commit()
    await db.refresh(current_user)

    # Return updated status (reuse the GET logic)
    return await get_onboarding_status(db=db, current_user=current_user)


# ---------------------------------------------------------------------------
# Welcome wizard endpoints (Phase 2)
#
# These persist Step 1/2/3 progress for the post-signup welcome wizard.
# Mounted on /users/me/* (the parent router prefix is /users) so the wizard
# can run before email verification and during trial.
# ---------------------------------------------------------------------------


@router.patch("/me/onboarding-step", response_model=OnboardingStepResponse)
async def patch_onboarding_step(
    body: OnboardingStepRequest,
    db: Annotated[AsyncSession, Depends(get_admin_db)],
    current_user: Annotated[User, Depends(get_current_active_user)],
) -> OnboardingStepResponse:
    """Persist welcome-wizard progress for the current user.

    Contract:
    - step=1 + complete writes accounts.name, accounts.team_size_bucket,
      users.role_at_signup, then sets users.onboarding_step_completed=1.
    - step=2 + complete writes accounts.primary_psa, then sets
      users.onboarding_step_completed=2.
    - step=3 + complete just sets users.onboarding_step_completed=3
      (invites are POSTed separately).
    - action="skip" ignores `data` entirely and only advances the step.
    - The new step must be >= current onboarding_step_completed (None=>0);
      otherwise 400. Idempotent re-PATCH of the same step succeeds.
    """
    current_step = current_user.onboarding_step_completed or 0
    if body.step < current_step:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail={
                "error": "step_cannot_decrease",
                "current_step": current_step,
                "requested_step": body.step,
            },
        )

    if body.action == "complete" and body.data is not None and body.step in (1, 2):
        # Load the user's account for field writes. Step 3 has no data writes.
        account_result = await db.execute(
            select(Account).where(Account.id == current_user.account_id)
        )
        account = account_result.scalar_one_or_none()
        if account is None:
            # Should never happen — user is required to have an account_id.
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="account_not_found",
            )

        if body.step == 1:
            data = body.data
            if data.company_name is not None:
                account.name = data.company_name
            if data.team_size_bucket is not None:
                account.team_size_bucket = data.team_size_bucket
            if data.role_at_signup is not None:
                current_user.role_at_signup = data.role_at_signup
        elif body.step == 2:
            data = body.data
            if data.primary_psa is not None:
                account.primary_psa = data.primary_psa

    current_user.onboarding_step_completed = body.step
    await db.commit()
    await db.refresh(current_user)

    return OnboardingStepResponse(
        onboarding_step_completed=current_user.onboarding_step_completed,
        onboarding_dismissed=current_user.onboarding_dismissed,
    )


@router.post("/me/onboarding-dismiss-rest", response_model=OnboardingStepResponse)
async def dismiss_onboarding_rest(
    db: Annotated[AsyncSession, Depends(get_admin_db)],
    current_user: Annotated[User, Depends(get_current_active_user)],
) -> OnboardingStepResponse:
    """Set users.onboarding_dismissed=TRUE — backs the wizard's "Skip the rest" button.

    Returns the same shape as the step PATCH so the frontend can update its
    local store from a single response.
    """
    current_user.onboarding_dismissed = True
    await db.commit()
    await db.refresh(current_user)

    return OnboardingStepResponse(
        onboarding_step_completed=current_user.onboarding_step_completed,
        onboarding_dismissed=current_user.onboarding_dismissed,
    )