SurfSense/surfsense_backend/app/routes/composio_routes.py

"""
Composio Connector OAuth Routes.

Handles OAuth flow for Composio-based integrations (Google Drive, Gmail, Calendar, etc.).
This provides a single connector that can connect to any Composio toolkit.

Endpoints:
- GET /composio/toolkits - List available Composio toolkits
- GET /auth/composio/connector/add - Initiate OAuth for a specific toolkit
- GET /auth/composio/connector/callback - Handle OAuth callback
- GET /connectors/{connector_id}/composio-drive/folders - List folders/files for Composio Google Drive
"""

import logging
from uuid import UUID

from fastapi import APIRouter, Depends, HTTPException, Query, Request
from fastapi.responses import RedirectResponse
from pydantic import ValidationError
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy.orm.attributes import flag_modified

from app.config import config
from app.db import (
    SearchSourceConnector,
    SearchSourceConnectorType,
    User,
    get_async_session,
)
from app.services.composio_service import (
    COMPOSIO_TOOLKIT_NAMES,
    INDEXABLE_TOOLKITS,
    TOOLKIT_TO_CONNECTOR_TYPE,
    ComposioService,
)
from app.users import current_active_user
from app.utils.connector_naming import (
    count_connectors_of_type,
    get_base_name_for_type,
)
from app.utils.oauth_security import OAuthStateManager

logger = logging.getLogger(__name__)

router = APIRouter()

# Map toolkit_id to frontend connector ID
TOOLKIT_TO_FRONTEND_CONNECTOR_ID = {
    "googledrive": "composio-googledrive",
    "gmail": "composio-gmail",
    "googlecalendar": "composio-googlecalendar",
}

# Initialize security utilities
_state_manager = None


def get_state_manager() -> OAuthStateManager:
    """Get or create OAuth state manager instance."""
    global _state_manager
    if _state_manager is None:
        if not config.SECRET_KEY:
            raise ValueError("SECRET_KEY must be set for OAuth security")
        _state_manager = OAuthStateManager(config.SECRET_KEY)
    return _state_manager


@router.get("/composio/toolkits")
async def list_composio_toolkits(user: User = Depends(current_active_user)):
    """
    List available Composio toolkits.

    Returns:
        JSON with list of available toolkits and their metadata.
    """
    if not ComposioService.is_enabled():
        raise HTTPException(
            status_code=503,
            detail="Composio integration is not enabled. Set COMPOSIO_ENABLED=TRUE and provide COMPOSIO_API_KEY.",
        )

    try:
        service = ComposioService()
        toolkits = service.list_available_toolkits()
        return {"toolkits": toolkits}
    except Exception as e:
        logger.error(f"Failed to list Composio toolkits: {e!s}")
        raise HTTPException(
            status_code=500, detail=f"Failed to list toolkits: {e!s}"
        ) from e


@router.get("/auth/composio/connector/add")
async def initiate_composio_auth(
    space_id: int,
    toolkit_id: str = Query(
        ..., description="Composio toolkit ID (e.g., 'googledrive', 'gmail')"
    ),
    user: User = Depends(current_active_user),
):
    """
    Initiate Composio OAuth flow for a specific toolkit.

    Query params:
        space_id: Search space ID to add connector to
        toolkit_id: Composio toolkit ID (e.g., "googledrive", "gmail", "googlecalendar")

    Returns:
        JSON with auth_url to redirect user to Composio authorization
    """
    if not ComposioService.is_enabled():
        raise HTTPException(
            status_code=503,
            detail="Composio integration is not enabled.",
        )

    if not space_id:
        raise HTTPException(status_code=400, detail="space_id is required")

    if toolkit_id not in COMPOSIO_TOOLKIT_NAMES:
        raise HTTPException(
            status_code=400,
            detail=f"Unknown toolkit: {toolkit_id}. Available: {list(COMPOSIO_TOOLKIT_NAMES.keys())}",
        )

    if not config.SECRET_KEY:
        raise HTTPException(
            status_code=500, detail="SECRET_KEY not configured for OAuth security."
        )

    try:
        # Generate secure state parameter with HMAC signature
        state_manager = get_state_manager()
        state_encoded = state_manager.generate_secure_state(
            space_id, user.id, toolkit_id=toolkit_id
        )

        # Build callback URL
        callback_url = config.COMPOSIO_REDIRECT_URI
        if not callback_url:
            # Fallback: construct from BACKEND_URL
            backend_url = config.BACKEND_URL or "http://localhost:8000"
            callback_url = f"{backend_url}/api/v1/auth/composio/connector/callback"

        # Initiate Composio OAuth
        service = ComposioService()
        # Use user.id as the entity ID in Composio (converted to string for Composio)
        entity_id = f"surfsense_{user.id}"

        connection_result = await service.initiate_connection(
            user_id=entity_id,
            toolkit_id=toolkit_id,
            redirect_uri=f"{callback_url}?state={state_encoded}",
        )

        auth_url = connection_result.get("redirect_url")
        if not auth_url:
            raise HTTPException(
                status_code=500, detail="Failed to get authorization URL from Composio"
            )

        logger.info(
            f"Initiating Composio OAuth for user {user.id}, toolkit {toolkit_id}, space {space_id}"
        )
        return {"auth_url": auth_url}

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Failed to initiate Composio OAuth: {e!s}", exc_info=True)
        raise HTTPException(
            status_code=500, detail=f"Failed to initiate Composio OAuth: {e!s}"
        ) from e


@router.get("/auth/composio/connector/callback")
async def composio_callback(
    request: Request,
    state: str | None = None,
    error: str | None = None,
    session: AsyncSession = Depends(get_async_session),
):
    """
    Handle Composio OAuth callback.

    Query params:
        state: Encoded state with space_id, user_id, and toolkit_id
        connected_account_id: Composio connected account ID (may not be present)
        error: OAuth error (if user denied access or error occurred)

    Returns:
        Redirect to frontend success page
    """
    try:
        # Handle OAuth errors
        if error:
            logger.warning(f"Composio OAuth error: {error}")
            space_id = None
            if state:
                try:
                    state_manager = get_state_manager()
                    data = state_manager.validate_state(state)
                    space_id = data.get("space_id")
                except Exception:
                    logger.warning("Failed to validate state in error handler")

            if space_id:
                return RedirectResponse(
                    url=f"{config.NEXT_FRONTEND_URL}/dashboard/{space_id}/connectors/callback?error=composio_oauth_denied"
                )
            else:
                return RedirectResponse(
                    url=f"{config.NEXT_FRONTEND_URL}/dashboard?error=composio_oauth_denied"
                )

        # Validate required parameters
        if not state:
            raise HTTPException(status_code=400, detail="Missing state parameter")

        # Validate and decode state with signature verification
        state_manager = get_state_manager()
        try:
            data = state_manager.validate_state(state)
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(
                status_code=400, detail=f"Invalid state parameter: {e!s}"
            ) from e

        user_id = UUID(data["user_id"])
        space_id = data["space_id"]
        toolkit_id = data.get("toolkit_id")

        if not toolkit_id:
            raise HTTPException(status_code=400, detail="Missing toolkit_id in state")

        toolkit_name = COMPOSIO_TOOLKIT_NAMES.get(toolkit_id, toolkit_id)

        logger.info(
            f"Processing Composio callback for user {user_id}, toolkit {toolkit_id}, space {space_id}"
        )

        # Initialize Composio service
        service = ComposioService()

        # Extract connected_account_id from query params (accepts both camelCase and snake_case)
        query_params = request.query_params
        final_connected_account_id = query_params.get(
            "connectedAccountId"
        ) or query_params.get("connected_account_id")

        # If we still don't have a connected_account_id, warn but continue
        # (the connector will be created but indexing won't work until updated)
        if not final_connected_account_id:
            logger.warning(
                f"Could not find connected_account_id for toolkit {toolkit_id}. "
                "The connector will be created but indexing may not work."
            )
        else:
            logger.info(
                f"Successfully got connected_account_id: {final_connected_account_id}"
            )
            # Wait for Composio to finish exchanging the auth code for tokens.
            try:
                service.wait_for_connection(final_connected_account_id, timeout=30.0)
            except Exception:
                logger.warning(
                    f"wait_for_connection timed out for {final_connected_account_id}, "
                    "proceeding anyway",
                    exc_info=True,
                )

        # Build entity_id for Composio API calls (same format as used in initiate)
        entity_id = f"surfsense_{user_id}"

        # Build connector config
        connector_config = {
            "composio_connected_account_id": final_connected_account_id,
            "toolkit_id": toolkit_id,
            "toolkit_name": toolkit_name,
            "is_indexable": toolkit_id in INDEXABLE_TOOLKITS,
        }

        # Get the specific connector type for this toolkit
        connector_type_str = TOOLKIT_TO_CONNECTOR_TYPE.get(toolkit_id)
        if not connector_type_str:
            raise HTTPException(
                status_code=400,
                detail=f"Unknown toolkit: {toolkit_id}. Available: {list(TOOLKIT_TO_CONNECTOR_TYPE.keys())}",
            )
        connector_type = SearchSourceConnectorType(connector_type_str)

        # Get the base name for this connector type (e.g., "Google Drive", "Gmail")
        base_name = get_base_name_for_type(connector_type)

        # FIRST: Get the email for this connected account
        # This is needed to determine if it's a reconnection (same email) or new account
        email = None
        try:
            email = await service.get_connected_account_email(
                connected_account_id=final_connected_account_id,
                entity_id=entity_id,
                toolkit_id=toolkit_id,
            )
            if email:
                logger.info(f"Retrieved email {email} for {toolkit_id} connector")
        except Exception as email_error:
            logger.warning(f"Could not get email for connector: {email_error!s}")

        # Generate the connector name (with email if available)
        # Format: "Gmail (Composio) - john@gmail.com" or "Gmail (Composio) 1" if no email
        if email:
            connector_name = f"{base_name} (Composio) - {email}"
        else:
            # Fallback to generic naming if email not available
            count = await count_connectors_of_type(
                session, connector_type, space_id, user_id
            )
            if count == 0:
                connector_name = f"{base_name} (Composio) 1"
            else:
                connector_name = f"{base_name} (Composio) {count + 1}"

        # Check if a connector with this SAME name already exists (reconnection case)
        # This allows multiple accounts (different emails) while supporting reconnection
        existing_connector_result = await session.execute(
            select(SearchSourceConnector).where(
                SearchSourceConnector.connector_type == connector_type,
                SearchSourceConnector.search_space_id == space_id,
                SearchSourceConnector.user_id == user_id,
                SearchSourceConnector.name == connector_name,
            )
        )
        existing_connector = existing_connector_result.scalars().first()

        if existing_connector:
            # This is a RECONNECTION of the same account - update existing connector
            old_connected_account_id = existing_connector.config.get(
                "composio_connected_account_id"
            )
            if (
                old_connected_account_id
                and old_connected_account_id != final_connected_account_id
            ):
                try:
                    deleted = await service.delete_connected_account(
                        old_connected_account_id
                    )
                    if deleted:
                        logger.info(
                            f"Deleted old Composio connected account {old_connected_account_id} "
                            f"before updating connector {existing_connector.id}"
                        )
                except Exception as delete_error:
                    logger.warning(
                        f"Error deleting old Composio connected account {old_connected_account_id}: {delete_error!s}"
                    )

            # Update existing connector with new connected_account_id
            # Merge new credentials with existing config to preserve user settings
            logger.info(
                f"Reconnecting existing Composio connector {existing_connector.id} ({connector_name}) "
                f"with new connected_account_id {final_connected_account_id}"
            )
            existing_config = (
                existing_connector.config.copy() if existing_connector.config else {}
            )
            existing_config.update(connector_config)
            existing_connector.config = existing_config

            flag_modified(existing_connector, "config")
            await session.commit()
            await session.refresh(existing_connector)

            frontend_connector_id = TOOLKIT_TO_FRONTEND_CONNECTOR_ID.get(
                toolkit_id, "composio-connector"
            )
            return RedirectResponse(
                url=f"{config.NEXT_FRONTEND_URL}/dashboard/{space_id}/connectors/callback?success=true&connector={frontend_connector_id}&connectorId={existing_connector.id}"
            )

        # This is a NEW account - create a new connector
        try:
            logger.info(f"Creating new Composio connector: {connector_name}")

            db_connector = SearchSourceConnector(
                name=connector_name,
                connector_type=connector_type,
                config=connector_config,
                search_space_id=space_id,
                user_id=user_id,
                is_indexable=toolkit_id in INDEXABLE_TOOLKITS,
            )

            session.add(db_connector)
            await session.commit()
            await session.refresh(db_connector)

            logger.info(
                f"Successfully created Composio connector {db_connector.id} for user {user_id}, toolkit {toolkit_id}"
            )

            # Get the frontend connector ID based on toolkit_id
            frontend_connector_id = TOOLKIT_TO_FRONTEND_CONNECTOR_ID.get(
                toolkit_id, "composio-connector"
            )
            return RedirectResponse(
                url=f"{config.NEXT_FRONTEND_URL}/dashboard/{space_id}/connectors/callback?success=true&connector={frontend_connector_id}&connectorId={db_connector.id}"
            )

        except IntegrityError as e:
            await session.rollback()
            logger.error(f"Database integrity error: {e!s}")
            raise HTTPException(
                status_code=409,
                detail=f"Database integrity error: {e!s}",
            ) from e
        except ValidationError as e:
            await session.rollback()
            logger.error(f"Validation error: {e!s}")
            raise HTTPException(
                status_code=400, detail=f"Invalid connector configuration: {e!s}"
            ) from e

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Unexpected error in Composio callback: {e!s}", exc_info=True)
        raise HTTPException(
            status_code=500, detail=f"Failed to complete Composio OAuth: {e!s}"
        ) from e


COMPOSIO_CONNECTOR_TYPES = {
    SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR,
    SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR,
    SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR,
}


@router.get("/auth/composio/connector/reauth")
async def reauth_composio_connector(
    space_id: int,
    connector_id: int,
    return_url: str | None = None,
    user: User = Depends(current_active_user),
    session: AsyncSession = Depends(get_async_session),
):
    """
    Initiate Composio re-authentication for an expired connected account.

    Uses Composio's refresh API so the same connected_account_id stays valid
    after the user completes the OAuth flow again.

    Query params:
        space_id: Search space ID the connector belongs to
        connector_id: ID of the existing Composio connector to re-authenticate
        return_url: Optional frontend path to redirect to after completion
    """
    if not ComposioService.is_enabled():
        raise HTTPException(
            status_code=503, detail="Composio integration is not enabled."
        )

    if not config.SECRET_KEY:
        raise HTTPException(
            status_code=500, detail="SECRET_KEY not configured for OAuth security."
        )

    try:
        result = await session.execute(
            select(SearchSourceConnector).filter(
                SearchSourceConnector.id == connector_id,
                SearchSourceConnector.user_id == user.id,
                SearchSourceConnector.search_space_id == space_id,
                SearchSourceConnector.connector_type.in_(COMPOSIO_CONNECTOR_TYPES),
            )
        )
        connector = result.scalars().first()
        if not connector:
            raise HTTPException(
                status_code=404,
                detail="Composio connector not found or access denied",
            )

        connected_account_id = connector.config.get("composio_connected_account_id")
        if not connected_account_id:
            raise HTTPException(
                status_code=400,
                detail="Composio connected account ID not found. Please reconnect the connector.",
            )

        # Build callback URL with secure state
        state_manager = get_state_manager()
        state_encoded = state_manager.generate_secure_state(
            space_id,
            user.id,
            toolkit_id=connector.config.get("toolkit_id", ""),
            connector_id=connector_id,
            return_url=return_url,
        )

        callback_base = config.COMPOSIO_REDIRECT_URI
        if not callback_base:
            backend_url = config.BACKEND_URL or "http://localhost:8000"
            callback_base = (
                f"{backend_url}/api/v1/auth/composio/connector/reauth/callback"
            )
        else:
            # Replace the normal callback path with the reauth one
            callback_base = callback_base.replace(
                "/auth/composio/connector/callback",
                "/auth/composio/connector/reauth/callback",
            )

        callback_url = f"{callback_base}?state={state_encoded}"

        service = ComposioService()
        refresh_result = service.refresh_connected_account(
            connected_account_id=connected_account_id,
            redirect_url=callback_url,
        )

        if refresh_result["redirect_url"] is None:
            # Token refreshed server-side; clear auth_expired immediately
            if connector.config.get("auth_expired"):
                connector.config = {**connector.config, "auth_expired": False}
                flag_modified(connector, "config")
                await session.commit()
            logger.info(
                f"Composio account {connected_account_id} refreshed server-side (no redirect needed)"
            )
            return {
                "success": True,
                "message": "Authentication refreshed successfully.",
            }

        logger.info(f"Initiating Composio re-auth for connector {connector_id}")
        return {"auth_url": refresh_result["redirect_url"]}

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Failed to initiate Composio re-auth: {e!s}", exc_info=True)
        raise HTTPException(
            status_code=500, detail=f"Failed to initiate Composio re-auth: {e!s}"
        ) from e


@router.get("/auth/composio/connector/reauth/callback")
async def composio_reauth_callback(
    request: Request,
    state: str | None = None,
    session: AsyncSession = Depends(get_async_session),
):
    """
    Handle Composio re-authentication callback.

    Clears the auth_expired flag and redirects the user back to the frontend.
    The connected_account_id has not changed — Composio refreshed it in place.
    """
    try:
        if not state:
            raise HTTPException(status_code=400, detail="Missing state parameter")

        state_manager = get_state_manager()
        try:
            data = state_manager.validate_state(state)
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(
                status_code=400, detail=f"Invalid state parameter: {e!s}"
            ) from e

        user_id = UUID(data["user_id"])
        space_id = data["space_id"]
        reauth_connector_id = data.get("connector_id")
        return_url = data.get("return_url")

        if not reauth_connector_id:
            raise HTTPException(status_code=400, detail="Missing connector_id in state")

        result = await session.execute(
            select(SearchSourceConnector).filter(
                SearchSourceConnector.id == reauth_connector_id,
                SearchSourceConnector.user_id == user_id,
                SearchSourceConnector.search_space_id == space_id,
            )
        )
        connector = result.scalars().first()
        if not connector:
            raise HTTPException(
                status_code=404,
                detail="Connector not found or access denied during re-auth callback",
            )

        # Wait for Composio to finish processing new tokens before proceeding.
        # Without this, get_access_token() may return stale credentials.
        connected_account_id = connector.config.get("composio_connected_account_id")
        if connected_account_id:
            try:
                service = ComposioService()
                service.wait_for_connection(connected_account_id, timeout=30.0)
            except Exception:
                logger.warning(
                    f"wait_for_connection timed out for connector {reauth_connector_id}, "
                    "proceeding anyway — tokens may not be ready yet",
                    exc_info=True,
                )

        # Clear auth_expired flag
        connector.config = {**connector.config, "auth_expired": False}
        flag_modified(connector, "config")
        await session.commit()
        await session.refresh(connector)

        logger.info(f"Composio re-auth completed for connector {reauth_connector_id}")

        if return_url and return_url.startswith("/"):
            return RedirectResponse(url=f"{config.NEXT_FRONTEND_URL}{return_url}")

        frontend_connector_id = TOOLKIT_TO_FRONTEND_CONNECTOR_ID.get(
            connector.config.get("toolkit_id", ""), "composio-connector"
        )
        return RedirectResponse(
            url=f"{config.NEXT_FRONTEND_URL}/dashboard/{space_id}/connectors/callback?success=true&connector={frontend_connector_id}&connectorId={reauth_connector_id}"
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in Composio reauth callback: {e!s}", exc_info=True)
        raise HTTPException(
            status_code=500, detail=f"Failed to complete Composio re-auth: {e!s}"
        ) from e


@router.get("/connectors/{connector_id}/composio-drive/folders")
async def list_composio_drive_folders(
    connector_id: int,
    parent_id: str | None = None,
    session: AsyncSession = Depends(get_async_session),
    user: User = Depends(current_active_user),
):
    """
    List folders AND files in user's Google Drive via Composio.

    Uses the same GoogleDriveClient / list_folder_contents path as the native
    connector, with Composio-sourced credentials.  This means auth errors
    propagate identically (Google returns 401 → exception → auth_expired flag).
    """
    from app.connectors.google_drive import GoogleDriveClient, list_folder_contents
    from app.utils.google_credentials import build_composio_credentials

    if not ComposioService.is_enabled():
        raise HTTPException(
            status_code=503,
            detail="Composio integration is not enabled.",
        )

    connector = None
    try:
        result = await session.execute(
            select(SearchSourceConnector).filter(
                SearchSourceConnector.id == connector_id,
                SearchSourceConnector.user_id == user.id,
                SearchSourceConnector.connector_type
                == SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR,
            )
        )
        connector = result.scalars().first()

        if not connector:
            raise HTTPException(
                status_code=404,
                detail="Composio Google Drive connector not found or access denied",
            )

        composio_connected_account_id = connector.config.get(
            "composio_connected_account_id"
        )
        if not composio_connected_account_id:
            raise HTTPException(
                status_code=400,
                detail="Composio connected account not found. Please reconnect the connector.",
            )

        credentials = build_composio_credentials(composio_connected_account_id)
        drive_client = GoogleDriveClient(session, connector_id, credentials=credentials)

        items, error = await list_folder_contents(drive_client, parent_id=parent_id)

        if error:
            error_lower = error.lower()
            if (
                "401" in error
                or "invalid_grant" in error_lower
                or "token has been expired or revoked" in error_lower
                or "invalid credentials" in error_lower
                or "authentication failed" in error_lower
            ):
                try:
                    if connector and not connector.config.get("auth_expired"):
                        connector.config = {**connector.config, "auth_expired": True}
                        flag_modified(connector, "config")
                        await session.commit()
                        logger.info(
                            f"Marked Composio connector {connector_id} as auth_expired"
                        )
                except Exception:
                    logger.warning(
                        f"Failed to persist auth_expired for connector {connector_id}",
                        exc_info=True,
                    )
                raise HTTPException(
                    status_code=400,
                    detail="Google Drive authentication expired. Please re-authenticate.",
                )
            raise HTTPException(
                status_code=500, detail=f"Failed to list folder contents: {error}"
            )

        folder_count = sum(1 for item in items if item.get("isFolder", False))
        file_count = len(items) - folder_count

        logger.info(
            f"Listed {len(items)} total items ({folder_count} folders, {file_count} files) for Composio connector {connector_id}"
            + (f" in folder {parent_id}" if parent_id else " in ROOT")
        )

        return {"items": items}

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error listing Composio Drive contents: {e!s}", exc_info=True)
        error_lower = str(e).lower()
        if (
            "invalid_grant" in error_lower
            or "token has been expired or revoked" in error_lower
            or "invalid credentials" in error_lower
            or "authentication failed" in error_lower
            or "401" in str(e)
        ):
            try:
                if connector and not connector.config.get("auth_expired"):
                    connector.config = {**connector.config, "auth_expired": True}
                    flag_modified(connector, "config")
                    await session.commit()
                    logger.info(
                        f"Marked Composio connector {connector_id} as auth_expired"
                    )
            except Exception:
                logger.warning(
                    f"Failed to persist auth_expired for connector {connector_id}",
                    exc_info=True,
                )
            raise HTTPException(
                status_code=400,
                detail="Google Drive authentication expired. Please re-authenticate.",
            ) from e
        raise HTTPException(
            status_code=500, detail=f"Failed to list Drive contents: {e!s}"
        ) from e