notifications.py

"""
Notifications database module

Structure:
users/{uid}/fcm_tokens (subcollection)
  └── {device_key} (document)
      ├── token: "actual_token_value"
      ├── created_at: timestamp
      └── time_zone: "America/New_York"

users/{uid} always carries daily_summary_enabled and daily_summary_hour_local
once a time_zone or preference write has run (and after the one-time backfill).
Those two fields are the write-time form of the Python defaults True / 22.
"""

from google.cloud.firestore_v1.base_query import FieldFilter
from google.cloud import firestore
from google.cloud.firestore import DELETE_FIELD
from ._client import db
from .cache import get_memory_cache
from .firestore_index_registry import DAILY_SUMMARY_RECIPIENTS_QUERY
import logging
from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast

logger = logging.getLogger(__name__)


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


def save_token(uid: str, data: Dict[str, Any]) -> None:
    """
    Store token in subcollection with device key as document ID
    Structure: users/{uid}/fcm_tokens/{device_key}
    Also maintains time_zone in main user document for backward compatibility
    Migrates legacy fcm_token to subcollection
    """
    device_key = data.get('device_key', 'unknown_default')
    token = data.get('fcm_token')
    time_zone = data.get('time_zone')

    user_ref = db.collection('users').document(uid)

    # Step 1: Migrate legacy token if exists
    user_doc = user_ref.get()
    user_data: Dict[str, Any] = {}
    if getattr(user_doc, "exists", False):
        user_data = _typed_doc(user_doc)
        legacy_token = user_data.get('fcm_token')

        if legacy_token:
            # Check if legacy token already exists in subcollection
            existing_tokens: List[object] = [
                t for t in (_typed_doc(d).get('token') for d in user_ref.collection('fcm_tokens').stream())
            ]

            if legacy_token not in existing_tokens:
                # Migrate to unknown_default
                user_ref.collection('fcm_tokens').document('unknown_default').set(
                    {
                        'token': legacy_token,
                        'time_zone': user_data.get('time_zone'),
                        'created_at': firestore.SERVER_TIMESTAMP,
                    },
                    merge=True,
                )

            # Remove legacy field
            user_ref.update({'fcm_token': DELETE_FIELD})

    # Step 2: If new token has proper device_key, replace unknown_default
    if device_key != 'unknown_default':
        unknown_ref = user_ref.collection('fcm_tokens').document('unknown_default')
        unknown_doc = unknown_ref.get()
        if getattr(unknown_doc, "exists", False):
            unknown_token = _typed_doc(unknown_doc).get('token')
            # Only delete if it's the same token being migrated to proper device_key
            if unknown_token == token:
                unknown_ref.delete()

    # Step 3: Save new token to subcollection
    user_ref.collection('fcm_tokens').document(device_key).set(
        {'token': token, 'time_zone': time_zone, 'created_at': firestore.SERVER_TIMESTAMP}, merge=True
    )

    # time_zone (when provided) plus any absent daily-summary schedule fields.
    # A user doc that becomes eligible via time_zone must already carry both
    # schedule fields so the indexed recipient query can match them.
    schedule_patch: Dict[str, Any] = {}
    if time_zone:
        schedule_patch['time_zone'] = time_zone
    schedule_patch.update(daily_summary_schedule_defaults(user_data))
    if schedule_patch:
        user_ref.set(schedule_patch, merge=True)


def get_user_time_zone(uid: str) -> Optional[str]:
    """Get timezone from main user document"""
    user_ref = db.collection('users').document(uid).get()
    if getattr(user_ref, "exists", False):
        user_data = _typed_doc(user_ref)
        tz = user_data.get('time_zone')
        return str(tz) if tz is not None else None
    return None


def set_user_time_zone_if_missing(uid: str, time_zone: str) -> bool:
    """Write ``time_zone`` on the user document only when it has none. Returns True when it wrote.

    ``save_token`` above is otherwise the only writer, and it runs from the mobile app's FCM
    registration. A desktop-only owner never registers a token, so their document never carried
    the field — and ``get_users_for_daily_summary`` selects users *by* it, so the daily-summary
    cron never saw them. Mobile stays authoritative: a zone already present is never replaced here.
    """
    user_ref = db.collection('users').document(uid)
    user_doc = user_ref.get()
    user_data = _typed_doc(user_doc) if getattr(user_doc, "exists", False) else {}
    if user_data.get('time_zone'):
        return False
    user_ref.set({'time_zone': time_zone, **daily_summary_schedule_defaults(user_data)}, merge=True)
    return True


# **************************************
# *** Daily Summary Time Preferences ***
# **************************************

# Default: 22:00 local time (10 PM); enabled unless the user turned it off.
DEFAULT_DAILY_SUMMARY_HOUR_LOCAL = 22
DEFAULT_DAILY_SUMMARY_ENABLED = True


def daily_summary_schedule_defaults(user_data: Mapping[str, Any]) -> Dict[str, Any]:
    """Return write-time defaults for whichever schedule fields are absent.

    Present values, including explicit ``False`` and hour ``0``, are never
    included. Empty dict when both fields are already on the document.
    """
    patch: Dict[str, Any] = {}
    if 'daily_summary_enabled' not in user_data:
        patch['daily_summary_enabled'] = DEFAULT_DAILY_SUMMARY_ENABLED
    if 'daily_summary_hour_local' not in user_data:
        patch['daily_summary_hour_local'] = DEFAULT_DAILY_SUMMARY_HOUR_LOCAL
    return patch


def get_daily_summary_hour_local(uid: str) -> int | None:
    """Get user's preferred daily summary hour in local time. Returns None if not set."""
    user_ref = db.collection('users').document(uid).get()
    if getattr(user_ref, "exists", False):
        user_data = _typed_doc(user_ref)
        value = user_data.get('daily_summary_hour_local')
        return int(value) if isinstance(value, (int, float)) else None
    return None


def set_daily_summary_hour_local(uid: str, hour_local: int) -> bool:
    """
    Set user's preferred daily summary hour in local time.

    Args:
        uid: User ID
        hour_local: Hour in local timezone (0-23)

    Returns:
        True if successful
    """
    if not (0 <= hour_local <= 23):
        raise ValueError(f"Invalid hour: {hour_local}. Must be 0-23.")

    user_ref = db.collection('users').document(uid)
    user_doc = user_ref.get()
    user_data = _typed_doc(user_doc) if getattr(user_doc, "exists", False) else {}
    user_ref.set(
        {**daily_summary_schedule_defaults(user_data), 'daily_summary_hour_local': hour_local},
        merge=True,
    )
    return True


def get_daily_summary_enabled(uid: str) -> bool:
    """Check if daily summary is enabled for user. Enabled by default."""
    user_ref = db.collection('users').document(uid).get()
    if getattr(user_ref, "exists", False):
        user_data = _typed_doc(user_ref)
        return bool(user_data.get('daily_summary_enabled', True))
    return True


def set_daily_summary_enabled(uid: str, enabled: bool) -> bool:
    """Enable or disable daily summary for user."""
    user_ref = db.collection('users').document(uid)
    user_doc = user_ref.get()
    user_data = _typed_doc(user_doc) if getattr(user_doc, "exists", False) else {}
    user_ref.set(
        {**daily_summary_schedule_defaults(user_data), 'daily_summary_enabled': enabled},
        merge=True,
    )
    return True


# **************************************
# *** Mentor Notification Frequency ***
# **************************************

# Default: 0 (disabled by default, user must explicitly enable)
# Range: 0-5 where 0=disabled, 1=most selective, 5=most proactive
DEFAULT_MENTOR_NOTIFICATION_FREQUENCY = 0


def get_mentor_notification_frequency(uid: str) -> int:
    """
    Get user's mentor notification frequency preference.
    Returns 0-5 where:
    - 0 = disabled
    - 1 = ultra selective (least frequent)
    - 3 = balanced (default)
    - 5 = very proactive (most frequent)

    Uses in-memory cache (30s TTL) + field projection to avoid reading the full
    user doc every 1s per stream. (#5439 sub-task 2)
    """
    cache = get_memory_cache()

    def fetch() -> int:
        doc = db.collection('users').document(uid).get(field_paths=['mentor_notification_frequency'])
        if getattr(doc, "exists", False):
            data = _typed_doc(doc)
            value = data.get('mentor_notification_frequency', DEFAULT_MENTOR_NOTIFICATION_FREQUENCY)
            return int(value) if isinstance(value, (int, float)) else DEFAULT_MENTOR_NOTIFICATION_FREQUENCY
        return DEFAULT_MENTOR_NOTIFICATION_FREQUENCY

    return cache.get_or_fetch(f"mentor_frequency:{uid}", fetch, ttl=30)


def set_mentor_notification_frequency(uid: str, frequency: int) -> bool:
    """
    Set user's mentor notification frequency preference.

    Args:
        uid: User ID
        frequency: Notification frequency (0-5)

    Returns:
        True if successful

    Raises:
        ValueError if frequency is not in valid range
    """
    if not (0 <= frequency <= 5):
        raise ValueError(f"Invalid frequency: {frequency}. Must be 0-5.")

    user_ref = db.collection('users').document(uid)
    user_ref.set({'mentor_notification_frequency': frequency}, merge=True)
    # Invalidate local cache so this instance sees the update immediately
    get_memory_cache().delete(f"mentor_frequency:{uid}")
    return True


def get_all_tokens(uid: str) -> list[str]:
    """Get all device tokens for a user from subcollection and legacy field"""
    tokens: List[str] = []

    # Get tokens from new subcollection
    token_docs = db.collection('users').document(uid).collection('fcm_tokens').stream()
    for doc in token_docs:
        token_data = _typed_doc(doc)
        token_value = token_data.get('token')
        if token_value:
            tokens.append(str(token_value))

    # Get legacy token from main user document (backward compatibility)
    user_ref = db.collection('users').document(uid).get()
    if getattr(user_ref, "exists", False):
        user_data = _typed_doc(user_ref)
        legacy_token = user_data.get('fcm_token')
        if legacy_token and legacy_token not in tokens:
            tokens.append(str(legacy_token))

    return tokens


def remove_invalid_token(token: str) -> None:
    """Remove invalid token using collection group query (rare operation)"""
    # Query across ALL users' fcm_tokens subcollections
    query = db.collection_group('fcm_tokens').where(filter=FieldFilter('token', '==', token)).limit(1)

    for doc in query.stream():
        doc.reference.delete()
        return


def remove_bulk_tokens(tokens: list[str]) -> None:
    """Remove multiple invalid tokens efficiently using IN queries and batch deletes"""
    if not tokens:
        return

    # Firestore IN queries support up to 30 items
    chunk_size = 30
    token_chunks = [tokens[i : i + chunk_size] for i in range(0, len(tokens), chunk_size)]

    for chunk in token_chunks:
        # Query for all tokens in this chunk at once
        query = db.collection_group('fcm_tokens').where(filter=FieldFilter('token', 'in', chunk))

        # Batch delete for efficiency
        batch = db.batch()
        count = 0

        for doc in query.stream():
            batch.delete(doc.reference)
            count += 1

            # Firestore batch limit is 500 operations
            if count >= 500:
                batch.commit()
                batch = db.batch()
                count = 0

        # Commit remaining deletes
        if count > 0:
            batch.commit()


def get_users_token_in_timezones(timezones: list[str]) -> List[str]:
    return _get_users_in_timezones(timezones, 'fcm_token')


def get_users_id_in_timezones(timezones: list[str]) -> List[Union[str, Tuple[str, List[str], Any]]]:
    return _get_users_in_timezones(timezones, 'id')


def get_users_for_daily_summary(timezones: list[str], target_local_hour: int) -> List[Tuple[str, List[str], Any]]:
    """
    Get users who should receive daily summary notifications.

    This function queries users who:
    1. Are in one of the provided timezones (where it's currently target_local_hour)
    2. Have daily_summary_hour_local set to target_local_hour OR have no preference (uses default)
    3. Have daily_summary_enabled not explicitly set to False

    Args:
        timezones: List of IANA timezone names where it's currently target_local_hour
        target_local_hour: The local hour we're sending notifications for (0-23)

    Returns:
        List of (uid, [tokens], time_zone) tuples.
    """
    if not timezones:
        return []

    users: List[Tuple[str, List[str], Any]] = []

    # 'Where in' query only supports 30 or fewer items in list so we split in chunks
    timezone_chunks = [timezones[i : i + 30] for i in range(0, len(timezones), 30)]

    for chunk in timezone_chunks:
        chunk_users: List[Tuple[str, List[str], Any]] = []
        try:
            # Query users in these timezones
            query = db.collection('users').where(filter=FieldFilter('time_zone', 'in', chunk))

            for user_doc in query.stream():
                uid = str(user_doc.id)
                user_data = _typed_doc(user_doc)

                # Check if daily summary is enabled (default: True)
                if user_data.get('daily_summary_enabled') is False:
                    continue

                # Check if user's preferred hour matches target hour
                # If not set, use default (22 = 10 PM)
                user_hour = user_data.get('daily_summary_hour_local', DEFAULT_DAILY_SUMMARY_HOUR_LOCAL)
                if user_hour != target_local_hour:
                    continue

                # Collect tokens from subcollection
                tokens: List[str] = []
                token_docs = db.collection('users').document(uid).collection('fcm_tokens').stream()
                for token_doc in token_docs:
                    token_data = _typed_doc(token_doc)
                    token_value = token_data.get('token')
                    if token_value:
                        tokens.append(str(token_value))

                # Add legacy token if exists and not already in list
                legacy_token = user_data.get('fcm_token')
                if legacy_token and legacy_token not in tokens:
                    tokens.append(str(legacy_token))

                # Tokenless users still get a record written. Generation is not
                # push delivery: a desktop-only owner has no FCM token and must
                # not be dropped here.
                time_zone = user_data.get('time_zone')
                chunk_users.append((uid, tokens, time_zone))

        except Exception as e:
            logger.error(f"Error querying chunk for daily summary: {e}")
        users.extend(chunk_users)

    return users


def get_users_for_daily_summary_indexed(
    timezones: list[str], target_local_hour: int
) -> List[Tuple[str, List[str], Any]]:
    """Select daily-summary recipients with server-side equality filters.

    The Python defaults (enabled True, hour 22) are materialized onto the user
    doc at write time and by ``backfill_daily_summary_schedule_fields``, so
    Firestore ``==`` matches the same set the legacy scan used to keep after a
    full ``time_zone IN`` pass. A user without those fields is invisible here
    until the backfill or a later write fills them. Chunks of >30 zones, token
    collection (subcollection + legacy ``fcm_token``), tokenless users, and
    per-chunk try/except-log-and-continue match ``get_users_for_daily_summary``.
    """
    if not timezones:
        return []

    users: List[Tuple[str, List[str], Any]] = []
    timezone_chunks = [timezones[i : i + 30] for i in range(0, len(timezones), 30)]

    for chunk in timezone_chunks:
        chunk_users: List[Tuple[str, List[str], Any]] = []
        try:
            query = DAILY_SUMMARY_RECIPIENTS_QUERY.build(
                db.collection('users'),
                {'enabled': True, 'hour_local': target_local_hour, 'time_zones': chunk},
                field_filter_factory=FieldFilter,
            )
            for user_doc in query.stream():
                uid = str(user_doc.id)
                user_data = _typed_doc(user_doc)
                tokens: List[str] = []
                token_docs = db.collection('users').document(uid).collection('fcm_tokens').stream()
                for token_doc in token_docs:
                    token_data = _typed_doc(token_doc)
                    token_value = token_data.get('token')
                    if token_value:
                        tokens.append(str(token_value))
                legacy_token = user_data.get('fcm_token')
                if legacy_token and legacy_token not in tokens:
                    tokens.append(str(legacy_token))
                time_zone = user_data.get('time_zone')
                chunk_users.append((uid, tokens, time_zone))
        except Exception as e:
            logger.error(f"Error querying chunk for daily summary: {e}")
        users.extend(chunk_users)

    return users


def _get_users_in_timezones(timezones: list[str], filter: str) -> List[Any]:
    """Query main user documents by timezone, then get tokens from subcollection and legacy field"""
    users: List[Any] = []

    # 'Where in' query only supports 30 or fewer items in list so we split in chunks
    timezone_chunks = [timezones[i : i + 30] for i in range(0, len(timezones), 30)]

    for chunk in timezone_chunks:
        chunk_users: List[Any] = []
        try:
            # Query main user documents by time_zone
            query = db.collection('users').where(filter=FieldFilter('time_zone', 'in', chunk))

            for user_doc in query.stream():
                uid = str(user_doc.id)
                user_data = _typed_doc(user_doc)

                # Collect tokens from subcollection
                tokens: List[str] = []
                token_docs = db.collection('users').document(uid).collection('fcm_tokens').stream()
                for token_doc in token_docs:
                    token_data = _typed_doc(token_doc)
                    token_value = token_data.get('token')
                    if token_value:
                        tokens.append(str(token_value))

                # Add legacy token if exists and not already in list
                legacy_token = user_data.get('fcm_token')
                if legacy_token and legacy_token not in tokens:
                    tokens.append(str(legacy_token))

                # Skip users with no tokens
                if not tokens:
                    continue

                if filter == 'fcm_token':
                    # Return flat list of tokens
                    chunk_users.extend(tokens)
                else:
                    # Return list of (uid, [tokens], time_zone) tuples
                    time_zone = user_data.get('time_zone')
                    chunk_users.append((uid, tokens, time_zone))

        except Exception as e:
            logger.error(f"Error querying chunk {chunk}: {e}")
        users.extend(chunk_users)

    return users