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

import secrets
import string
from datetime import datetime, timezone
from typing import Annotated

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

from app.api.endpoints.auth import _mint_session_tokens
from app.core.admin_database import get_admin_db
from app.core.config import settings
from app.models.account import Account
from app.models.account_invite import AccountInvite
from app.models.oauth_identity import OAuthIdentity
from app.models.user import User
from app.schemas.oauth import OAuthCallbackPayload, OAuthCallbackResponse
from app.services.billing import BillingService
from app.services.oauth_providers import (
    google_exchange_code,
    microsoft_exchange_code,
    OAuthProfile,
)

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


def _generate_display_code(length: int = 8) -> str:
    """Match the helper used by /auth/register — A-Z + 0-9, length 8."""
    alphabet = string.ascii_uppercase + string.digits
    return "".join(secrets.choice(alphabet) for _ in range(length))


async def _sign_in_or_register(
    db: AsyncSession,
    provider: str,
    profile: OAuthProfile,
    *,
    account_invite_code: str | None = None,
    invited_email: str | None = None,
) -> tuple[User, bool]:
    """Returns (user, is_new_user). Idempotent on (provider, provider_subject).

    When ``account_invite_code`` is supplied (from the /accept-invite flow),
    a brand-new user is created inside the invited account instead of getting
    a personal account + Pro trial. Mismatch between the OAuth profile email
    and ``invited_email`` raises ``invite_email_mismatch`` per the spec
    contract that mirrors the email+password register path.
    """
    identity = (
        await db.execute(
            select(OAuthIdentity).where(
                OAuthIdentity.provider == provider,
                OAuthIdentity.provider_subject == profile.provider_subject,
            )
        )
    ).scalar_one_or_none()

    if identity:
        user = (
            await db.execute(select(User).where(User.id == identity.user_id))
        ).scalar_one()
        return user, False

    user = (
        await db.execute(select(User).where(User.email == profile.email))
    ).scalar_one_or_none()
    is_new_user = user is None

    # If the user arrived via an invite link but already has a ResolutionFlow
    # account (e.g., previously signed up with email+password), silently
    # linking the OAuth identity to that existing account would bypass the
    # invite — they'd stay in their personal account and the invite would
    # never be consumed. Fail loud instead so they can sign in and accept the
    # invite from the dashboard. The "invited user wants to transfer accounts"
    # case is a v2 concern.
    if account_invite_code and not is_new_user:
        raise HTTPException(
            status_code=400,
            detail={
                "error": "email_already_registered_use_login",
                "message": (
                    "An account already exists for this email. Please sign in "
                    "instead, then accept the invite from your dashboard."
                ),
            },
        )

    invite_record: AccountInvite | None = None
    if is_new_user and account_invite_code:
        # SELECT FOR UPDATE so two concurrent OAuth callbacks can't both
        # consume the same invite code.
        invite_record = (
            await db.execute(
                select(AccountInvite)
                .where(AccountInvite.code == account_invite_code)
                .with_for_update()
            )
        ).scalar_one_or_none()
        if invite_record is None or not invite_record.is_valid:
            raise HTTPException(
                status_code=400,
                detail={"error": "invite_invalid_or_expired_or_revoked"},
            )
        # Verify the OAuth profile email matches what was invited. We compare
        # against the invite row directly (source of truth), but also accept
        # the client-supplied invited_email as a defensive equality check.
        if invite_record.email.lower() != profile.email.lower():
            raise HTTPException(
                status_code=400,
                detail={"error": "invite_email_mismatch"},
            )
        if invited_email and invited_email.lower() != invite_record.email.lower():
            raise HTTPException(
                status_code=400,
                detail={"error": "invite_email_mismatch"},
            )

    if is_new_user:
        if invite_record is not None:
            # Join the invited account directly — no personal account, no
            # trial creation.
            user = User(
                email=profile.email,
                name=profile.name,
                password_hash=None,
                account_id=invite_record.account_id,
                account_role=invite_record.role,
                role="engineer",
                email_verified_at=datetime.now(timezone.utc),
            )
            db.add(user)
            await db.flush()
            invite_record.accepted_by_id = user.id
            invite_record.used_at = datetime.now(timezone.utc)
            await db.flush()
        else:
            account = Account(
                name=f"{profile.name}'s Account",
                display_code=_generate_display_code(),
            )
            db.add(account)
            await db.flush()
            user = User(
                email=profile.email,
                name=profile.name,
                password_hash=None,
                account_id=account.id,
                account_role="owner",
                role="engineer",
                email_verified_at=datetime.now(timezone.utc),
            )
            db.add(user)
            await db.flush()
            account.owner_id = user.id
            await db.flush()
            # start_trial commits internally; flushed account/user above.
            await BillingService.start_trial(db, account.id)

    db.add(
        OAuthIdentity(
            user_id=user.id,
            provider=provider,
            provider_subject=profile.provider_subject,
            provider_email_at_link=profile.email,
        )
    )
    await db.commit()
    await db.refresh(user)
    return user, is_new_user


@router.post("/google/callback", response_model=OAuthCallbackResponse)
async def google_callback(
    payload: OAuthCallbackPayload,
    db: Annotated[AsyncSession, Depends(get_admin_db)],
) -> OAuthCallbackResponse:
    if not settings.GOOGLE_CLIENT_ID:
        raise HTTPException(status_code=503, detail="Google sign-in not configured")
    redirect_uri = f"{settings.OAUTH_REDIRECT_BASE}/auth/google/callback"
    profile = await google_exchange_code(payload.code, redirect_uri)
    user, is_new = await _sign_in_or_register(
        db,
        "google",
        profile,
        account_invite_code=payload.account_invite_code,
        invited_email=payload.invited_email,
    )
    token = await _mint_session_tokens(user, db)
    await db.commit()
    return OAuthCallbackResponse(
        access_token=token.access_token,
        refresh_token=token.refresh_token,
        is_new_user=is_new,
        idle_expires_at=token.idle_expires_at,
        absolute_expires_at=token.absolute_expires_at,
    )


@router.post("/microsoft/callback", response_model=OAuthCallbackResponse)
async def microsoft_callback(
    payload: OAuthCallbackPayload,
    db: Annotated[AsyncSession, Depends(get_admin_db)],
) -> OAuthCallbackResponse:
    if not settings.MS_CLIENT_ID:
        raise HTTPException(status_code=503, detail="Microsoft sign-in not configured")
    redirect_uri = f"{settings.OAUTH_REDIRECT_BASE}/auth/microsoft/callback"
    profile = await microsoft_exchange_code(payload.code, redirect_uri)
    user, is_new = await _sign_in_or_register(
        db,
        "microsoft",
        profile,
        account_invite_code=payload.account_invite_code,
        invited_email=payload.invited_email,
    )
    token = await _mint_session_tokens(user, db)
    await db.commit()
    return OAuthCallbackResponse(
        access_token=token.access_token,
        refresh_token=token.refresh_token,
        is_new_user=is_new,
        idle_expires_at=token.idle_expires_at,
        absolute_expires_at=token.absolute_expires_at,
    )