integrations.py

# async-blockers: no-import-scope
# async-blockers: no-changed-range-scope  # pre-existing patterns surfaced by type-annotation import changes
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from typing import Any, Callable, Dict, List, Literal, Optional, cast
from pydantic import BaseModel, Field
from urllib.parse import urlencode
import os
import secrets
import json
import base64
import hashlib
from datetime import datetime, timezone
import httpx

import database.users as users_db
import database.redis_db as redis_db
from utils.other import endpoints as auth
from utils.log_sanitizer import sanitize
from utils.llm.gateway_error_contract import BYOK_RATE_LIMIT_ERROR_DETAIL, is_byok_rate_limit_gateway_error
from utils.subscription import is_trial_paywalled
from utils.executors import run_blocking, db_executor, llm_executor
from utils.integrations_registry import oauth_authorization_query, resolve_integration_provider
from utils.retrieval.tools.google_utils import (
    GMAIL_READONLY_SCOPE,
    GOOGLE_INTEGRATION_KEY,
    google_integration_has_scope,
)
import logging

logger = logging.getLogger(__name__)

router = APIRouter()
_auth_module = cast(Any, auth)

# OAuth state management
OAUTH_STATE_EXPIRY = 600  # 10 minutes

http_client: Optional[httpx.AsyncClient] = None

# Templates
templates = Jinja2Templates(directory="templates")


def get_http_client() -> httpx.AsyncClient:
    """Get or create the HTTP client instance."""
    global http_client
    if http_client is None:
        http_client = httpx.AsyncClient(timeout=10.0)
    return http_client


async def close_http_client():
    """Close the HTTP client and cleanup resources."""
    global http_client
    if http_client is not None:
        await http_client.aclose()
        http_client = None


def render_oauth_response(
    request: Request,
    app_key: str,
    success: bool = True,
    redirect_url: Optional[str] = None,
    error_type: Optional[str] = None,
) -> HTMLResponse:
    """
    Render OAuth callback response using template.

    Args:
        request: FastAPI request object
        app_key: Integration app key (google_calendar, whoop)
        success: Whether the OAuth flow was successful
        redirect_url: Deep link URL to redirect to (for success case)
        error_type: Type of error (missing_code, invalid_state, config_error, server_error)
    """
    resolved = resolve_integration_provider(app_key)
    config = resolved[1] if resolved else {'name': app_key.title()}

    if success:
        context: Dict[str, Any] = {
            'title': f"{config['name']} Auth",
            'icon': '✓',
            'message': 'Authentication Successful!',
            'description': 'Redirecting back to Omi...',
            'redirect_url': redirect_url or f'omi://{app_key}/callback?error=unknown',
            'show_spinner': True,
        }
    else:
        error_messages: Dict[str, str] = {
            'missing_code': 'No authorization code received from {}.'.format(config['name']),
            'invalid_state': 'Invalid or expired authentication request.',
            'config_error': '{} OAuth not properly configured.'.format(config['name']),
            'server_error': 'An error occurred during authentication.',
        }

        context = {
            'title': f"{config['name']} Auth Error",
            'icon': '❌',
            'message': f"{'Security' if error_type == 'invalid_state' else 'Configuration' if error_type == 'config_error' else 'Authentication'} Error",
            'description': error_messages.get(error_type or 'unknown', 'An error occurred.'),
            'redirect_url': f'omi://{app_key}/callback?error={error_type or "unknown"}',
            'show_spinner': False,
        }

    return templates.TemplateResponse(request, 'oauth_callback.html', context)


def validate_and_consume_oauth_state(state_token: Optional[str]) -> Optional[Dict[str, str]]:
    """
    Validate OAuth state token and return associated data.
    Deletes the state token after validation to prevent replay attacks.

    Returns:
        Dict with 'uid' and 'app_key' if valid, None if invalid/expired
    """
    if not state_token:
        return None

    state_key = f"oauth_state:{state_token}"
    # Atomic get-and-delete: an OAuth state is single-use, so consuming it must be one operation.
    # A separate GET then DELETE lets two concurrent callbacks carrying the same state both read the
    # value before either delete runs, which weakens replay protection -- and offloading the consume
    # to the db_executor thread pool makes that interleaving reachable. GETDEL removes it atomically,
    # so only one caller ever receives the value.
    state_data_str = redis_db.r.getdel(state_key)

    if not state_data_str:
        return None

    try:
        loaded: object = json.loads(state_data_str.decode() if isinstance(state_data_str, bytes) else state_data_str)
        state_data = cast(Dict[str, str], loaded) if isinstance(loaded, dict) else {}
        # GETDEL above already removed the key atomically; no separate delete needed.
        return state_data
    except Exception as e:
        logger.error(f"Error parsing state data: {e}")
        return None


# Request/Response models
class IntegrationData(BaseModel):
    """Data for an integration connection"""

    connected: bool = True
    access_token: Optional[str] = None
    refresh_token: Optional[str] = None


class AppleHealthSyncData(BaseModel):
    """Health data synced from Apple Health on iOS device"""

    period_days: int = Field(default=7, description="Number of days of data")

    # Steps data
    total_steps: Optional[int] = Field(default=None, description="Total steps in period")
    average_steps_per_day: Optional[float] = Field(default=None, description="Average steps per day")
    daily_steps: Optional[List[Dict[str, Any]]] = Field(
        default=None, description="Daily steps breakdown [{date, steps}]"
    )

    # Sleep data
    total_sleep_hours: Optional[float] = Field(default=None, description="Total sleep hours")
    total_in_bed_hours: Optional[float] = Field(default=None, description="Total time in bed hours")
    sleep_sessions_count: Optional[int] = Field(default=None, description="Number of sleep sessions")
    sleep_sessions: Optional[List[Dict[str, Any]]] = Field(default=None, description="Sleep session details")
    daily_sleep: Optional[List[Dict[str, Any]]] = Field(
        default=None, description="Daily sleep breakdown [{date, sleepHours}]"
    )

    # Heart rate data
    heart_rate_average: Optional[float] = Field(default=None, description="Average heart rate")
    heart_rate_min: Optional[float] = Field(default=None, description="Minimum heart rate")
    heart_rate_max: Optional[float] = Field(default=None, description="Maximum heart rate")

    # Active energy data
    total_active_energy: Optional[float] = Field(default=None, description="Total active energy kcal")
    average_active_energy_per_day: Optional[float] = Field(default=None, description="Average daily active energy")
    daily_active_energy: Optional[List[Dict[str, Any]]] = Field(
        default=None, description="Daily energy breakdown [{date, calories}]"
    )

    # Workouts data
    workouts: Optional[List[Dict[str, Any]]] = Field(default=None, description="List of workout records")


class IntegrationResponse(BaseModel):
    """Response containing integration status"""

    connected: bool = Field(description="Whether the integration is connected")
    app_key: str = Field(description="Integration app key")


class IntegrationMutationResponse(BaseModel):
    status: str
    app_key: str


class AppleHealthSyncResponse(BaseModel):
    status: str
    app_key: str
    synced_at: str
    data_types_synced: list[str] = Field(default_factory=list)


# *****************************
# ********** ROUTES ***********
# *****************************


# Integrations that are not stored under their own key: they are derived from another
# grant, and each requires a specific scope on that grant.
DERIVED_INTEGRATIONS = {
    'gmail': (GOOGLE_INTEGRATION_KEY, GMAIL_READONLY_SCOPE),
}


class ConnectorSynthesisRequest(BaseModel):
    model_config = {"extra": "forbid"}

    source: Literal['calendar', 'gmail', 'notes']
    items: List[str] = Field(..., min_length=1, max_length=200)
    existing_memories: List[str] = Field(default_factory=list, max_length=200)


class ConnectorSynthesisTask(BaseModel):
    model_config = {"extra": "forbid"}

    description: str
    priority: str = "medium"
    due_at: str = ""


class ConnectorSynthesisResponse(BaseModel):
    model_config = {"extra": "forbid"}

    memories: List[str]
    tasks: List[ConnectorSynthesisTask]
    profile: str = ""


@router.post("/v1/connectors/synthesize", tags=['integrations'], response_model=ConnectorSynthesisResponse)
async def synthesize_connector_data(
    body: ConnectorSynthesisRequest,
    uid: str = Depends(
        cast(Callable[..., str], _auth_module.with_rate_limit(auth.get_current_user_uid, "connectors:synthesize"))
    ),
):
    """Return-only calendar/gmail/notes synthesis through the managed memories feature.

    Does not write Firestore. Desktop connector importers call this instead of building
    their own prompts and inventing memories via Anthropic Haiku chat completions, then
    persist through the normal memory/task write APIs.
    """
    from utils.llm import connector_synthesis

    if await run_blocking(db_executor, is_trial_paywalled, uid, 'desktop'):
        raise HTTPException(status_code=402, detail='trial_expired')
    try:
        synthesis = await run_blocking(
            llm_executor,
            lambda: connector_synthesis.synthesize_connector_items(
                uid,
                body.source,
                body.items,
                existing_memories=body.existing_memories,
            ),
        )
    except Exception as e:
        if not is_byok_rate_limit_gateway_error(e):
            raise
        logger.warning('Connector synthesis halted because the configured BYOK provider is rate limited')
        raise HTTPException(status_code=429, detail=BYOK_RATE_LIMIT_ERROR_DETAIL) from None
    if synthesis is None:
        raise HTTPException(status_code=502, detail="connector_synthesis_failed")
    return ConnectorSynthesisResponse(
        memories=list(synthesis.memories),
        tasks=[
            ConnectorSynthesisTask(description=t.description, priority=t.priority, due_at=t.due_at)
            for t in synthesis.tasks
        ],
        profile=synthesis.profile or "",
    )


@router.get("/v1/integrations/{app_key}", response_model=IntegrationResponse, tags=['integrations'])
def get_integration(app_key: str, uid: str = Depends(auth.get_current_user_uid)):
    """Get integration connection status for the current user.

    Gmail has no grant of its own — it rides the Google Calendar OAuth grant and is
    connected only when that grant actually carries the Gmail scope.
    """
    if app_key in DERIVED_INTEGRATIONS:
        source_key, required_scope = DERIVED_INTEGRATIONS[app_key]
        source = users_db.get_integration(uid, source_key)
        connected = bool(source and source.get('connected')) and google_integration_has_scope(source, required_scope)
        return IntegrationResponse(connected=connected, app_key=app_key)

    integration = users_db.get_integration(uid, app_key)

    if integration and integration.get('connected'):
        return IntegrationResponse(connected=True, app_key=app_key)
    else:
        return IntegrationResponse(connected=False, app_key=app_key)


@router.put("/v1/integrations/{app_key}", tags=['integrations'], response_model=IntegrationMutationResponse)
def save_integration(app_key: str, data: IntegrationData, uid: str = Depends(auth.get_current_user_uid)):
    """Save or update an integration connection."""
    # Convert Pydantic model to dict, excluding None values
    integration_data = data.model_dump(exclude_none=True)

    users_db.set_integration(uid, app_key, integration_data)

    return {"status": "ok", "app_key": app_key}


@router.delete("/v1/integrations/{app_key}", status_code=204, tags=['integrations'])
def delete_integration(app_key: str, uid: str = Depends(auth.get_current_user_uid)):
    """Delete an integration connection.

    Deleting a derived integration deletes the grant it rides on — there is no
    separate token to revoke.
    """
    if app_key in DERIVED_INTEGRATIONS:
        app_key = DERIVED_INTEGRATIONS[app_key][0]

    success = users_db.delete_integration(uid, app_key)

    if not success:
        raise HTTPException(status_code=404, detail="Integration not found")

    return None


@router.put("/v1/integrations/apple-health/sync", response_model=AppleHealthSyncResponse, tags=['integrations'])
def sync_apple_health_data(data: AppleHealthSyncData, uid: str = Depends(auth.get_current_user_uid)):
    """
    Sync Apple Health data from the iOS device.

    This endpoint receives health data collected from Apple HealthKit on the user's
    iPhone/Apple Watch and stores it for use in chat queries.

    Unlike other integrations that use OAuth, Apple Health data is pushed from the device.
    """
    # Build the health data structure
    health_data: Dict[str, Any] = {
        'period_days': data.period_days,
    }

    # Steps
    if data.total_steps is not None:
        health_data['steps'] = {
            'total': data.total_steps,
            'average_per_day': data.average_steps_per_day or (data.total_steps / max(data.period_days, 1)),
            'period_days': data.period_days,
            'daily': data.daily_steps or [],  # Daily breakdown [{date, steps}]
        }

    # Sleep
    if data.total_sleep_hours is not None or data.sleep_sessions:
        health_data['sleep'] = {
            'total_sleep_hours': data.total_sleep_hours or 0,
            'total_in_bed_hours': data.total_in_bed_hours or 0,
            'sessions_count': data.sleep_sessions_count or 0,
            'sessions': data.sleep_sessions or [],
            'daily': data.daily_sleep or [],  # Daily breakdown [{date, sleepHours}]
        }

    # Heart rate
    if data.heart_rate_average is not None:
        health_data['heart_rate'] = {
            'average': data.heart_rate_average,
            'minimum': data.heart_rate_min,
            'maximum': data.heart_rate_max,
        }

    # Active energy
    if data.total_active_energy is not None:
        health_data['active_energy'] = {
            'total': data.total_active_energy,
            'average_per_day': data.average_active_energy_per_day
            or (data.total_active_energy / max(data.period_days, 1)),
            'daily': data.daily_active_energy or [],  # Daily breakdown [{date, calories}]
        }

    # Workouts
    if data.workouts:
        health_data['workouts'] = data.workouts

    # Save the integration with health data
    integration_data: Dict[str, Any] = {
        'connected': True,
        'health_data': health_data,
        'last_synced': datetime.now(timezone.utc).isoformat(),
    }

    users_db.set_integration(uid, 'apple_health', integration_data)

    return {
        "status": "ok",
        "app_key": "apple_health",
        "synced_at": integration_data['last_synced'],
        "data_types_synced": list(health_data.keys()),
    }


# *****************************
# ****** OAuth Initiation *****
# *****************************


class OAuthUrlResponse(BaseModel):
    """Response containing OAuth authorization URL"""

    auth_url: str = Field(description="OAuth authorization URL to open in browser")


@router.get("/v1/integrations/{app_key}/oauth-url", response_model=OAuthUrlResponse, tags=['integrations'])
def get_oauth_url(app_key: str, uid: str = Depends(auth.get_current_user_uid)):
    """
    Get OAuth authorization URL for an integration.
    Frontend opens this URL in browser to start OAuth flow.
    Uses secure random state tokens to prevent CSRF attacks.

    A derived integration (Gmail) authorizes through the grant it rides on, so the
    whole flow — state, provider config and callback — runs under the source key.
    """
    if app_key in DERIVED_INTEGRATIONS:
        app_key = DERIVED_INTEGRATIONS[app_key][0]

    base_url = os.getenv('BASE_API_URL')
    if not base_url:
        logger.error(f'ERROR: BASE_API_URL not configured for integration OAuth')
        raise HTTPException(status_code=500, detail="BASE_API_URL not configured")

    resolved = resolve_integration_provider(app_key)
    if not resolved or resolved[1]['kind'] != 'oauth':
        raise HTTPException(status_code=400, detail=f"Unsupported integration: {app_key}")
    provider_key, provider = resolved
    oauth = cast(Dict[str, Any], provider['oauth'])

    # Generate cryptographically secure random state token
    state_token = secrets.token_urlsafe(32)

    # Store state mapping in Redis with expiry
    try:
        state_key = f"oauth_state:{state_token}"
        state_data = {'uid': uid, 'app_key': provider_key, 'created_at': datetime.now(timezone.utc).isoformat()}
        redis_db.r.setex(state_key, OAUTH_STATE_EXPIRY, json.dumps(state_data))
    except Exception as e:
        logger.error(f'ERROR: Failed to store OAuth state in Redis: {e}')
        raise HTTPException(status_code=500, detail=f"Failed to initialize OAuth flow: {str(e)}")

    client_id_env = cast(str, oauth['client_id_env'])
    client_id = os.getenv(client_id_env)
    if not client_id:
        logger.error(f"ERROR: {client_id_env} not configured for {provider['name']} integration OAuth")
        raise HTTPException(status_code=500, detail=f"{provider['name']} not configured - {client_id_env} missing")

    base_url_clean = base_url.rstrip('/')
    redirect_uri = f"{base_url_clean}{oauth['redirect_path']}"

    params: Dict[str, str] = {
        'client_id': client_id,
        'redirect_uri': redirect_uri,
        'state': state_token,
    }
    params.update(oauth_authorization_query(provider))

    if oauth.get('requires_pkce'):
        code_verifier = secrets.token_urlsafe(32)
        code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).decode().rstrip('=')
        verifier_key = f"oauth_code_verifier:{state_token}"
        redis_db.r.setex(verifier_key, OAUTH_STATE_EXPIRY, code_verifier)
        params['code_challenge'] = code_challenge

    auth_url = f"{oauth['auth_base']}?{urlencode(params)}"
    logger.info(f"Generated {provider['name']} OAuth URL for user {uid}")

    return OAuthUrlResponse(auth_url=auth_url)


# *****************************
# ******* OAuth Callbacks *****
# *****************************


class OAuthProviderConfig(BaseModel):
    """Configuration for OAuth provider-specific logic"""

    token_endpoint: str
    token_request_type: str = "form"
    token_request_data: Dict[str, Any]
    additional_headers: Dict[str, str] = {}

    async def fetch_additional_data(self, client: httpx.AsyncClient, access_token: str) -> Dict[str, Any]:
        """Hook for fetching provider-specific data after token exchange"""
        return {}


async def handle_oauth_callback(
    request: Request,
    app_key: str,
    code: Optional[str],
    state: Optional[str],
    provider_config: OAuthProviderConfig,
) -> HTMLResponse:
    """
    Generic OAuth callback handler that works for all providers.

    Args:
        request: FastAPI request object
        app_key: Integration app key (google_calendar, whoop)
        code: Authorization code from OAuth provider
        state: State token for CSRF protection
        provider_config: Provider-specific configuration

    Returns:
        HTMLResponse with OAuth callback page
    """
    if not code or not state:
        return render_oauth_response(request, app_key, success=False, error_type='missing_code')

    # Validate state token
    state_data = await run_blocking(db_executor, validate_and_consume_oauth_state, state)
    if not state_data or state_data.get('app_key') != app_key:
        return render_oauth_response(request, app_key, success=False, error_type='invalid_state')

    uid = state_data.get('uid')
    if not uid:
        return render_oauth_response(request, app_key, success=False, error_type='invalid_state')

    try:
        client = get_http_client()

        if provider_config.token_request_type == "form":
            token_response = await client.post(
                provider_config.token_endpoint,
                headers={
                    'Content-Type': 'application/x-www-form-urlencoded',
                    **provider_config.additional_headers,
                },
                data=provider_config.token_request_data,
            )
        elif provider_config.token_request_type == "json":
            token_response = await client.post(
                provider_config.token_endpoint,
                headers={
                    'Content-Type': 'application/json',
                    **provider_config.additional_headers,
                },
                json=provider_config.token_request_data,
            )
        else:  # params
            token_response = await client.post(
                provider_config.token_endpoint,
                params=provider_config.token_request_data,
                headers=provider_config.additional_headers,
            )

        if token_response.status_code == 200:
            token_data = token_response.json()
            access_token = token_data.get('access_token', '')
            refresh_token = token_data.get('refresh_token')

            if not access_token:
                logger.info(f'{app_key}: No access token received in response')
                return render_oauth_response(request, app_key, success=False, error_type='server_error')

            integration_data: Dict[str, Any] = {
                'connected': True,
                'access_token': access_token,
                # Google returns only the scopes the user actually approved; store them so
                # scope-derived integrations (Gmail) can tell granted from merely requested.
                'granted_scopes': (token_data.get('scope') or '').split(),
            }

            if refresh_token:
                integration_data['refresh_token'] = refresh_token

            try:
                additional_data = await provider_config.fetch_additional_data(client, access_token)
                integration_data.update(additional_data)
            except Exception as e:
                logger.error(f'{app_key}: Error fetching additional data: {e}')

            # Store in Firebase
            try:
                await run_blocking(db_executor, users_db.set_integration, uid, app_key, integration_data)
            except Exception as e:
                logger.error(f'{app_key}: Error storing tokens in Firebase: {e}')
                return render_oauth_response(request, app_key, success=False, error_type='server_error')

            deep_link = f'omi://{app_key}/callback?success=true'

            return render_oauth_response(request, app_key, success=True, redirect_url=deep_link)
        else:
            error_body = token_response.text[:500] if token_response.text else "No error body"
            logger.error(f'{app_key}: Token exchange failed with HTTP {token_response.status_code}')
            logger.error(f'{app_key}: Error response: {sanitize(error_body)}')
            return render_oauth_response(request, app_key, success=False, error_type='server_error')

    except Exception as e:
        logger.error(f'{app_key}: Unexpected error during OAuth callback: {e}')
        return render_oauth_response(request, app_key, success=False, error_type='server_error')


@router.get(
    '/v2/integrations/{app_key}/callback',
    response_class=HTMLResponse,
    tags=['integrations', 'oauth'],
)
async def oauth_callback(
    request: Request,
    app_key: str,
    code: Optional[str] = Query(None),
    state: Optional[str] = Query(None),
):
    resolved = resolve_integration_provider(app_key)
    if not resolved or resolved[1]['kind'] != 'oauth':
        return render_oauth_response(request, app_key, success=False, error_type='config_error')
    provider_key, provider = resolved
    oauth = cast(Dict[str, Any], provider['oauth'])

    client_id_env = cast(str, oauth['client_id_env'])
    client_secret_env = cast(str, oauth['client_secret_env'])
    client_id = os.getenv(client_id_env)
    client_secret = os.getenv(client_secret_env)
    base_url = os.getenv('BASE_API_URL')

    if not client_id or not client_secret or not base_url:
        return render_oauth_response(request, provider_key, success=False, error_type='config_error')

    base_url_clean = base_url.rstrip('/')
    redirect_uri = f"{base_url_clean}{oauth['redirect_path']}"

    config = OAuthProviderConfig(
        token_endpoint=cast(str, oauth['token_endpoint']),
        token_request_type=cast(str, oauth.get('token_request_type', 'form')),
        token_request_data={
            'code': code,
            'client_id': client_id,
            'client_secret': client_secret,
            'redirect_uri': redirect_uri,
            'grant_type': 'authorization_code',
        },
    )
    return await handle_oauth_callback(request, provider_key, code, state, config)


@router.on_event("shutdown")  # type: ignore[reportDeprecated]  # FastAPI on_event still functional; lifespan migration would change app wiring
async def shutdown_http_client():
    """Cleanup HTTP client on app shutdown."""
    await close_http_client()