headquarter/apps/api/src/api/auth.py

import logging
from secrets import token_urlsafe
from typing import Any, AsyncGenerator, Literal, cast

import httpx
from fastapi import APIRouter, Cookie, Depends, HTTPException, Response, status
from fastapi.responses import RedirectResponse
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from src.auth.cookies import build_cookie_options
from src.auth.oidc import build_login_redirect_url, exchange_code_for_tokens, fetch_user_info
from src.auth.session import create_session_cookie, decode_session_cookie
from src.config import Settings
from src.database import SessionLocal
from src.models.user import User

logger = logging.getLogger(__name__)

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


async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
    async with SessionLocal() as session:
        yield session


@router.get(
    "/login",
    summary="Initiate OAuth login",
    description="Redirects to the configured OAuth provider (Authentik) to start the authentication flow.",
    response_class=RedirectResponse,
)
async def login(next: str = "/") -> RedirectResponse:
    """Initiate OAuth2 login flow.

    Args:
        next: URL to redirect to after successful authentication.

    Returns:
        RedirectResponse to the OAuth provider's authorization endpoint.
    """
    settings = Settings()
    redirect_uri = f"{settings.api_base_url}/auth/callback"
    state = token_urlsafe(24)
    location = build_login_redirect_url(
        settings=settings,
        redirect_uri=redirect_uri,
        state=state,
    )
    logger.debug("Auth login initiated: redirect_uri=%s, next=%s", redirect_uri, next)
    response = RedirectResponse(location)
    response.set_cookie("auth_state", state, httponly=True, samesite="lax")
    response.set_cookie("auth_next", next, httponly=True, samesite="lax")
    return response


@router.get("/callback")
async def callback(
    code: str,
    state: str,
    auth_state: str | None = Cookie(default=None),
    auth_next: str | None = Cookie(default="/"),
    session: AsyncSession = Depends(get_db_session),
) -> RedirectResponse:
    logger.debug("Auth callback received: code=%s... state=%s", code[:10] if code else "None", state[:10] if state else "None")
    
    if auth_state is None or auth_state != state:
        logger.warning("State mismatch: cookie=%s, param=%s", auth_state, state)
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="invalid state")

    settings = Settings()
    redirect_uri = f"{settings.api_base_url}/auth/callback"
    logger.debug("Exchanging code for tokens (redirect_uri=%s)", redirect_uri)
    
    async with httpx.AsyncClient() as client:
        try:
            token_payload = await exchange_code_for_tokens(
                settings=settings,
                code=code,
                redirect_uri=redirect_uri,
                client=client,
            )
            logger.info("Token exchange successful")
        except Exception as exc:
            logger.error("Token exchange failed: %s", exc)
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"token exchange failed: {exc}")
        
        try:
            user_info = await fetch_user_info(
                settings=settings,
                access_token=token_payload["access_token"],
                client=client,
            )
            logger.debug("User info fetched successfully")
        except Exception as exc:
            logger.error("User info fetch failed: %s", exc)
            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="failed to fetch user info")

    authentik_id = str(user_info.get("sub", ""))
    email = str(user_info.get("email", f"{authentik_id}@authentik.local"))
    name = str(user_info.get("name", email))
    logger.debug("User info: authentik_id=%s, email=%s, name=%s", authentik_id, email, name)

    try:
        user = await session.scalar(select(User).where(User.authentik_id == authentik_id))
        if user is None:
            logger.debug("Creating new user: authentik_id=%s", authentik_id)
            user = User(email=email, name=name, authentik_id=authentik_id, avatar_url=None)
            session.add(user)
            await session.commit()
            await session.refresh(user)
            logger.info("New user created: id=%s", user.id)
        else:
            logger.debug("Existing user found: id=%s, updating info", user.id)
            user.email = email
            user.name = name
            await session.commit()
    except Exception as exc:
        logger.error("Database error during user lookup/creation: %s", exc)
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="database error")

    # Create session cookie
    session_cookie = create_session_cookie(settings=settings, user_id=str(user.id))
    
    cookie_options = build_cookie_options(settings)
    cookie_samesite = cast(Literal["lax", "strict", "none"], cookie_options["samesite"])
    cookie_secure = bool(cookie_options["secure"])
    cookie_domain = str(cookie_options["domain"]) if cookie_options.get("domain") else None
    
    logger.info("Auth callback complete for user id=%s, redirecting to %s", user.id, auth_next)
    
    # Redirect to frontend with the original next path
    redirect_url = f"{settings.web_base_url}{auth_next}"
    redirect_response = RedirectResponse(url=redirect_url)
    
    redirect_response.set_cookie(
        "session", 
        session_cookie, 
        httponly=True, 
        samesite=cookie_samesite, 
        secure=cookie_secure,
        domain=cookie_domain,
    )
    redirect_response.delete_cookie("auth_state", samesite="lax", domain=cookie_domain)
    redirect_response.delete_cookie("auth_next", samesite="lax", domain=cookie_domain)
    
    return redirect_response


@router.post("/logout")
async def logout(response: Response) -> dict[str, str]:
    settings = Settings()
    cookie_options = build_cookie_options(settings)
    cookie_samesite = cast(Literal["lax", "strict", "none"], cookie_options["samesite"])
    cookie_secure = bool(cookie_options["secure"])
    cookie_domain = str(cookie_options["domain"]) if cookie_options.get("domain") else None

    response.delete_cookie("session", samesite=cookie_samesite, secure=cookie_secure, domain=cookie_domain)
    return {"status": "ok"}


@router.get("/me")
async def me(
    session_cookie: str | None = Cookie(default=None, alias="session"),
    session: AsyncSession = Depends(get_db_session),
) -> dict[str, Any]:
    logger.debug("Auth /me called, cookie present: %s", bool(session_cookie))
    
    if not session_cookie:
        logger.warning("Auth /me: missing session cookie")
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="missing session")

    settings = Settings()
    logger.debug("Auth /me: cookie_domain=%s, cookie_secure=%s, cookie_samesite=%s", 
                settings.cookie_domain, settings.cookie_secure, settings.cookie_samesite)
    
    try:
        payload = decode_session_cookie(settings=settings, cookie_value=session_cookie)
        user_id = payload["user_id"]
        logger.debug("Auth /me: decoded session for user_id=%s", user_id)
    except ValueError as exc:
        logger.warning("Auth /me: invalid session: %s", exc)
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=str(exc))

    user = await session.get(User, user_id)
    if user is None:
        logger.warning("Auth /me: user not found for id=%s", user_id)
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="user not found")

    logger.info("Auth /me: success for user=%s", user.email)
    return {
        "user": {
            "id": str(user.id),
            "email": user.email,
            "name": user.name,
            "avatar_url": user.avatar_url or "",
        }
    }