apps.py

from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, cast

from google.cloud.firestore_v1.base_query import BaseCompositeFilter, FieldFilter
from google.cloud.firestore import ArrayUnion, ArrayRemove

from ulid import ULID

from models.app import App, UsageHistoryType
from .redis_db import get_generic_cache, set_generic_cache
from ._client import db
import logging

logger = logging.getLogger(__name__)

# Shared with utils.apps (list + invalidation). Keep every reader and the invalidation path on this
# one constant: a second literal is how a cache ends up populated but never cleared.
PUBLIC_APPROVED_APPS_CACHE_KEY = 'get_public_approved_apps_data'

# BaseCompositeFilter expects Operator enum but accepts 'AND' string at runtime.
# Typed as Any to satisfy pyright without importing StructuredQuery (which fails
# on some google-cloud-firestore versions).
_AND_OP: Any = 'AND'


def _typed_doc(doc: Any) -> Dict[str, Any]:
    raw: object = doc.to_dict()
    return cast(Dict[str, Any], raw) if isinstance(raw, dict) else {}


# *****************************
# ********** CRUD *************
# *****************************

apps_collection = 'plugins_data'
app_analytics_collection = 'plugins'
testers_collection = 'testers'


def get_app_by_id_db(app_id: str) -> Optional[Dict[str, Any]]:
    app_ref = db.collection(apps_collection).document(app_id)
    doc = app_ref.get()
    if doc.exists:
        raw: object = doc.to_dict()
        return cast(Dict[str, Any], raw) if isinstance(raw, dict) else None
    return None


def get_audio_apps_count(app_ids: List[str]) -> int:
    if not app_ids or len(app_ids) == 0:
        return 0
    filters = [FieldFilter('id', 'in', app_ids), FieldFilter('external_integration.triggers_on', '==', 'audio_bytes')]
    apps_ref = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).count().get()
    return apps_ref[0][0].value


def get_private_apps_db(uid: str) -> List[Dict[str, Any]]:
    filters = [FieldFilter('uid', '==', uid), FieldFilter('private', '==', True)]
    private_apps = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).stream()
    data = [_typed_doc(doc) for doc in private_apps]
    return data


# This returns public unapproved apps of all users
def get_unapproved_public_apps_db() -> List[Dict[str, Any]]:
    filters = [FieldFilter('approved', '==', False), FieldFilter('private', '==', False)]
    public_apps = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).stream()
    return [_typed_doc(doc) for doc in public_apps]


def get_public_approved_apps_db() -> List[Dict[str, Any]]:
    filters = [FieldFilter('approved', '==', True), FieldFilter('private', '==', False)]
    public_apps = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).stream()
    return [_typed_doc(doc) for doc in public_apps]


def get_public_approved_apps_cached_db() -> List[Dict[str, Any]]:
    """The approved+public app set, read through the marketplace's shared 10-minute Redis cache.

    Same key, TTL, reduction and invalidation as `utils.apps.get_approved_available_apps`, so a
    reader here can never serve a staler view than the list the user just came from.
    """
    cached = get_generic_cache(PUBLIC_APPROVED_APPS_CACHE_KEY)
    if cached:
        return cast(List[Dict[str, Any]], cached)
    reduced = [App.reduce_dict(app) for app in get_public_approved_apps_db()]
    set_generic_cache(PUBLIC_APPROVED_APPS_CACHE_KEY, reduced, 60 * 10)  # 10 minutes cached
    return reduced


def get_popular_apps_db() -> List[Dict[str, Any]]:
    filters = [FieldFilter('approved', '==', True), FieldFilter('is_popular', '==', True)]
    popular_apps = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).stream()
    return [_typed_doc(doc) for doc in popular_apps]


def set_app_popular_db(app_id: str, popular: bool) -> None:
    app_ref = db.collection(apps_collection).document(app_id)
    app_ref.update({'is_popular': popular})


def search_apps_db(
    uid: str,
    category: str | None = None,
    capability: str | None = None,
    my_apps: bool = False,
    installed_apps: bool = False,
    enabled_app_ids: List[str] | None = None,
) -> List[Dict[str, Any]]:
    """
    Optimized search function that applies filters at database level.
    Uses smart filter ordering to minimize data fetched from Firestore.

    Note: Rating filter is NOT applied here as rating_avg is calculated from Redis,
    not stored in Firestore. Apply rating filter after fetching from DB.

    Args:
        uid: User ID for private apps and filtering
        category: Filter by category ID
        capability: Filter by capability ID
        my_apps: Only return user's own apps
        installed_apps: Only return user's enabled apps
        enabled_app_ids: Pre-fetched list of enabled app IDs (for installed_apps filter)

    Returns:
        List of app dictionaries matching the filters
    """
    filters: List[FieldFilter] = []
    # Whether the primary read is the whole approved+public app set. That set is 3k+ documents and
    # streaming it per request is what made `?q=` search a p50-13s / p90-30s endpoint in prod; the
    # marketplace list path already serves the same documents from Redis, so read through it here too.
    reads_public_set = False

    # 1. Apply most restrictive filter first
    if my_apps:
        filters.append(FieldFilter('uid', '==', uid))

    elif installed_apps:
        if not enabled_app_ids or len(enabled_app_ids) == 0:
            # User has no enabled apps
            return []

        if len(enabled_app_ids) > 30:
            # Firestore 'in' limited to 30 items
            # Query public approved apps first, then add user's own apps
            reads_public_set = True
        else:
            # Query by specific IDs
            filters.append(FieldFilter('id', 'in', enabled_app_ids))

    else:
        # Default: Public approved apps
        reads_public_set = True

    # 2. Add category filter
    if category and not my_apps:  # Don't add if already filtering by my_apps
        filters.append(FieldFilter('category', '==', category))

    # 3. Add capability filter
    if capability and not my_apps:
        filters.append(FieldFilter('capabilities', 'array_contains', capability))

    # Execute query with all filters
    apps: List[Dict[str, Any]] = []
    if reads_public_set:
        apps = get_public_approved_apps_cached_db()
        # category/capability were server-side filters on the Firestore read this replaces; my_apps is
        # False on this branch, so both are unconditional here.
        if category:
            apps = [app for app in apps if app.get('category') == category]
        if capability:
            apps = [app for app in apps if capability in (app.get('capabilities') or [])]
    elif filters:
        query = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters))
        apps = [_typed_doc(doc) for doc in query.stream()]

    # For installed_apps with > 30 enabled apps, we need to also fetch user's own apps
    # because the main query only returns approved+public apps
    if installed_apps and enabled_app_ids and len(enabled_app_ids) > 30:
        enabled_set = set(enabled_app_ids)
        # Filter to only enabled apps from the public approved set
        apps = [app for app in apps if app.get('id') in enabled_set]

        # Also fetch user's own enabled apps (which may be private or unapproved)
        user_apps_filter = FieldFilter('uid', '==', uid)
        user_apps_query = db.collection(apps_collection).where(filter=user_apps_filter)
        user_apps = [_typed_doc(doc) for doc in user_apps_query.stream()]

        # Add user's own enabled apps that aren't already in the list
        existing_ids = {app.get('id') for app in apps}
        for user_app in user_apps:
            if user_app.get('id') in enabled_set and user_app.get('id') not in existing_ids:
                apps.append(user_app)

    # Post-filter for category if my_apps is enabled
    if my_apps and category:
        apps = [app for app in apps if app.get('category') == category]

    # Post-filter for capability if my_apps is enabled
    if my_apps and capability:
        apps = [app for app in apps if capability in app.get('capabilities', [])]

    return apps


# This returns public unapproved apps for a user
def get_public_unapproved_apps_db(uid: str) -> List[Dict[str, Any]]:
    filters = [FieldFilter('approved', '==', False), FieldFilter('uid', '==', uid), FieldFilter('private', '==', False)]
    public_apps = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).stream()
    return [_typed_doc(doc) for doc in public_apps]


def get_apps_for_tester_db(uid: str) -> List[Dict[str, Any]]:
    tester_ref = db.collection(testers_collection).document(uid)
    doc = tester_ref.get()
    if doc.exists:
        apps = _typed_doc(doc).get('apps', [])
        if not apps:
            return []
        filters = [FieldFilter('approved', '==', False), FieldFilter('id', 'in', apps)]
        public_apps = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).stream()
        return [_typed_doc(doc) for doc in public_apps]
    return []


def add_app_to_db(app_data: Dict[str, Any]) -> None:
    app_ref = db.collection(apps_collection)
    app_ref.add(app_data, app_data['id'])


def upsert_app_to_db(app_data: Dict[str, Any]) -> None:
    app_ref = db.collection(apps_collection).document(app_data['id'])
    app_ref.set(app_data)


def update_app_in_db(app_data: Dict[str, Any]) -> None:
    app_ref = db.collection(apps_collection).document(app_data['id'])
    app_ref.update(app_data)


def delete_app_from_db(app_id: str) -> None:
    app_ref = db.collection(apps_collection).document(app_id)
    app_ref.delete()


def update_app_visibility_in_db(app_id: str, private: bool) -> None:
    app_ref = db.collection(apps_collection).document(app_id)
    if 'private' in app_id and not private:
        app = _typed_doc(app_ref.get())
        if not app:
            # The private app document is gone (deleted, or a stale read-cache pointed the caller
            # here). There is nothing to republish, so skip the delete-and-recreate instead of
            # dereferencing None below (which raised TypeError -> 500).
            return
        app_ref.delete()
        new_app_id = app_id.split('-private')[0] + '-' + str(ULID())
        app['id'] = new_app_id
        app['private'] = private
        app_ref = db.collection(apps_collection).document(new_app_id)
        app_ref.set(app)
    else:
        app_ref.update({'private': private})


def change_app_approval_status(app_id: str, approved: bool) -> None:
    app_ref = db.collection(apps_collection).document(app_id)
    app_ref.update({'approved': approved, 'status': 'approved' if approved else 'rejected'})


def get_app_usage_history_db(app_id: str) -> List[Dict[str, Any]]:
    usage = db.collection(app_analytics_collection).document(app_id).collection('usage_history').stream()
    return [_typed_doc(doc) for doc in usage]


def get_app_memory_created_integration_usage_count_db(app_id: str) -> Any:
    usage = (
        db.collection(app_analytics_collection)
        .document(app_id)
        .collection('usage_history')
        .where(filter=FieldFilter('type', '==', UsageHistoryType.memory_created_external_integration))
        .count()
        .get()
    )
    return usage[0][0].value


def get_app_memory_prompt_usage_count_db(app_id: str) -> Any:
    usage = (
        db.collection(app_analytics_collection)
        .document(app_id)
        .collection('usage_history')
        .where(filter=FieldFilter('type', '==', UsageHistoryType.memory_created_prompt))
        .count()
        .get()
    )
    return usage[0][0].value


def get_app_chat_message_sent_usage_count_db(app_id: str) -> Any:
    usage = (
        db.collection(app_analytics_collection)
        .document(app_id)
        .collection('usage_history')
        .where(filter=FieldFilter('type', '==', UsageHistoryType.chat_message_sent))
        .count()
        .get()
    )
    return usage[0][0].value


def get_app_usage_count_db(app_id: str) -> Any:
    usage = db.collection(app_analytics_collection).document(app_id).collection('usage_history').count().get()
    return usage[0][0].value


# ********************************
# *********** REVIEWS ************
# ********************************


def set_app_review_in_db(app_id: str, uid: str, review: Dict[str, Any]) -> None:
    app_ref = db.collection(apps_collection).document(app_id).collection('reviews').document(uid)
    app_ref.set(review)


# ********************************
# ************ TESTER ************
# ********************************


def add_tester_db(data: Dict[str, Any]) -> None:
    app_ref = db.collection(testers_collection).document(data['uid'])
    app_ref.set(data)


def add_app_access_for_tester_db(app_id: str, uid: str) -> None:
    app_ref = db.collection(testers_collection).document(uid)
    app_ref.update({'apps': ArrayUnion([app_id])})


def remove_app_access_for_tester_db(app_id: str, uid: str) -> None:
    app_ref = db.collection(testers_collection).document(uid)
    app_ref.update({'apps': ArrayRemove([app_id])})


def remove_tester_db(uid: str) -> None:
    app_ref = db.collection(testers_collection).document(uid)
    app_ref.delete()


def can_tester_access_app_db(app_id: str, uid: str) -> bool:
    app_ref = db.collection(testers_collection).document(uid)
    doc = app_ref.get()
    if doc.exists:
        return app_id in _typed_doc(doc).get('apps', [])
    return False


def is_tester_db(uid: str) -> bool:
    app_ref = db.collection(testers_collection).document(uid)
    return app_ref.get().exists


# ********************************
# *********** APPS USAGE *********
# ********************************


def record_app_usage(
    uid: str,
    app_id: str,
    usage_type: UsageHistoryType,
    conversation_id: Optional[str] = None,
    message_id: Optional[str] = None,
    timestamp: Optional[datetime] = None,
) -> Dict[str, Any]:
    if not conversation_id and not message_id:
        raise ValueError('memory_id or message_id must be provided')

    data: Dict[str, Any] = {
        'uid': uid,
        'memory_id': conversation_id,
        'message_id': message_id,
        'timestamp': datetime.now(timezone.utc) if timestamp is None else timestamp,
        'type': usage_type,
    }

    db.collection(app_analytics_collection).document(app_id).collection('usage_history').document(
        conversation_id or message_id
    ).set(data)
    return data


# ********************************
# *********** PERSONAS ***********
# ********************************


def delete_persona_db(persona_id: str) -> None:
    persona_ref = db.collection(apps_collection).document(persona_id)
    persona_ref.delete()


def get_personas_by_username_db(persona_id: str) -> Optional[List[Dict[str, Any]]]:
    persona_ref = db.collection(apps_collection).where('username', '==', persona_id)
    docs = persona_ref.get()
    if not docs:
        return None
    return [{**_typed_doc(doc), 'doc_id': doc.id} for doc in docs]


def get_persona_by_username_db(username: str) -> Optional[Dict[str, Any]]:
    filters = [FieldFilter('username', '==', username), FieldFilter('capabilities', 'array_contains', 'persona')]
    persona_ref = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).limit(1)
    docs = persona_ref.get()
    if not docs:
        return None
    doc = next(iter(docs), None)
    if not doc:
        return None
    raw: object = doc.to_dict()
    return cast(Dict[str, Any], raw) if isinstance(raw, dict) else None


def get_persona_by_id_db(persona_id: str) -> Optional[Dict[str, Any]]:
    persona_ref = db.collection(apps_collection).document(persona_id)
    doc = persona_ref.get()
    if doc.exists:
        raw: object = doc.to_dict()
        return cast(Dict[str, Any], raw) if isinstance(raw, dict) else None
    return None


def get_persona_by_uid_db(uid: str) -> Optional[Dict[str, Any]]:
    filters = [FieldFilter('uid', '==', uid), FieldFilter('capabilities', 'array_contains', 'persona')]
    persona_ref = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).limit(1)
    docs = persona_ref.get()
    if not docs:
        return None
    doc = next(iter(docs), None)
    if not doc:
        return None
    raw: object = doc.to_dict()
    return cast(Dict[str, Any], raw) if isinstance(raw, dict) else None


def get_user_persona_by_uid(uid: str) -> Optional[Dict[str, Any]]:
    filters = [
        FieldFilter('capabilities', 'array_contains', 'persona'),
        FieldFilter('category', '==', 'personality-emulation'),
        FieldFilter('uid', '==', uid),
    ]
    persona_ref = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).limit(1)
    docs = persona_ref.get()
    if not docs:
        return None
    doc = next(iter(docs), None)
    if not doc:
        return None
    return {'id': doc.id, **_typed_doc(doc)}


def get_persona_by_twitter_handle_db(handle: str) -> Optional[Dict[str, Any]]:
    filters = [FieldFilter('category', '==', 'personality-emulation'), FieldFilter('twitter.username', '==', handle)]
    persona_ref = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).limit(1)
    docs = persona_ref.get()
    if not docs:
        return None
    doc = next(iter(docs), None)
    if not doc:
        return None
    return {'id': doc.id, **_typed_doc(doc)}


def get_persona_by_username_twitter_handle_db(username: str, handle: str) -> Optional[Dict[str, Any]]:
    filters = [
        FieldFilter('username', '==', username),
        FieldFilter('category', '==', 'personality-emulation'),
        FieldFilter('twitter.username', '==', handle),
    ]
    persona_ref = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).limit(1)
    docs = persona_ref.get()
    if not docs:
        return None
    doc = next(iter(docs), None)
    if not doc:
        return None
    return {'id': doc.id, **_typed_doc(doc)}


def get_omi_personas_by_uid_db(uid: str) -> List[Dict[str, Any]]:
    filters = [FieldFilter('uid', '==', uid), FieldFilter('capabilities', 'array_contains', 'persona')]
    persona_ref = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters))
    docs = persona_ref.get()
    if not docs:
        return []
    typed_docs = [_typed_doc(doc) for doc in docs]
    docs_out = [d for d in typed_docs if 'omi' in d.get('connected_accounts', [])]
    return docs_out


def get_omi_persona_apps_by_uid_db(uid: str) -> List[Dict[str, Any]]:
    filters = [FieldFilter('uid', '==', uid), FieldFilter('category', '==', 'personality-emulation')]
    persona_ref = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters))
    docs = persona_ref.get()
    if not docs:
        return []
    docs_out = [_typed_doc(doc) for doc in docs]
    return docs_out


def update_persona_in_db(persona_data: Dict[str, Any]) -> None:
    persona_ref = db.collection(apps_collection).document(persona_data['id'])
    persona_ref.update(persona_data)


def migrate_app_owner_id_db(new_id: str, old_id: str) -> None:
    filters = [FieldFilter('uid', '==', old_id)]
    apps_ref = db.collection(apps_collection).where(filter=BaseCompositeFilter(_AND_OP, filters)).stream()
    for app in apps_ref:
        app_ref = db.collection(apps_collection).document(app.id)
        app_ref.update({'uid': new_id})


def create_api_key_db(app_id: str, api_key_data: Dict[str, Any]) -> Dict[str, Any]:
    """Create a new API key for an app in the database"""
    api_key_ref = db.collection(apps_collection).document(app_id).collection('api_keys').document(api_key_data['id'])
    api_key_ref.set(api_key_data)
    return api_key_data


def get_api_key_by_hash_db(app_id: str, hashed_key: str) -> Optional[Dict[str, Any]]:
    """Get an API key by its hash value"""
    filters = [FieldFilter('hashed', '==', hashed_key)]
    api_keys_ref = (
        db.collection(apps_collection)
        .document(app_id)
        .collection('api_keys')
        .where(filter=BaseCompositeFilter(_AND_OP, filters))
        .limit(1)
    )
    docs = api_keys_ref.get()
    if not docs:
        return None
    doc = next(iter(docs), None)
    if not doc:
        return None
    raw: object = doc.to_dict()
    return cast(Dict[str, Any], raw) if isinstance(raw, dict) else None


def list_api_keys_db(app_id: str) -> List[Dict[str, Any]]:
    """List all API keys for an app (excluding the hashed values)"""
    api_keys_ref = (
        db.collection(apps_collection)
        .document(app_id)
        .collection('api_keys')
        .order_by('created_at', direction='DESCENDING')
        .stream()
    )
    return [{k: v for k, v in _typed_doc(doc).items() if k != 'hashed'} for doc in api_keys_ref]


def delete_api_key_db(app_id: str, key_id: str) -> bool:
    """Delete an API key"""
    api_key_ref = db.collection(apps_collection).document(app_id).collection('api_keys').document(key_id)
    api_key_ref.delete()
    return True