bdbot/handlers/scheduler.py

"""Scheduler for daily birthday notifications."""
import telebot
from datetime import datetime, timedelta, date
from typing import Optional, Dict, List, Tuple
from collections import defaultdict
from apscheduler.schedulers.blocking import BlockingScheduler
from apscheduler.triggers.cron import CronTrigger
import pytz
from sqlalchemy.orm import Session
from database import get_db_session, User, Chat, UserChat
from messages import format_birthday_greeting, format_multiple_birthdays_greetings
from config import Config


def send_birthday_notifications(bot: telebot.TeleBot) -> None:
    """Send birthday notifications to all chats for users with birthday today."""
    db = get_db_session()
    try:
        today = datetime.now().date()
        day = today.day
        month = today.month
        
        # Get all users with birthday today
        users_with_birthday = db.query(User).filter(
            User.birthday_day == day,
            User.birthday_month == month
        ).all()
        
        if not users_with_birthday:
            return
        
        # Group users by chat: chat_id -> list of (first_name, theme, user_id)
        chat_users: Dict[int, List[Tuple[str, str, int]]] = defaultdict(list)
        
        # For each user, find all their chats
        for user in users_with_birthday:
            # Get all chats where user is a member
            user_chats = db.query(Chat).join(UserChat).filter(
                UserChat.user_id == user.user_id,
                Chat.bot_is_admin == True
            ).all()
            
            for chat in user_chats:
                try:
                    # Check if user is still in chat
                    member = bot.get_chat_member(chat.chat_id, user.user_id)
                    if member.status not in ['left', 'kicked']:
                        chat_users[chat.chat_id].append(
                            (user.first_name, user.preference_theme, user.user_id)
                        )
                except telebot.apihelper.ApiTelegramException as e:
                    # Handle errors: user blocked bot, bot removed from chat, etc.
                    if e.error_code == 403:
                        # Bot was blocked or removed from chat
                        # Update chat status
                        chat.bot_is_admin = False
                        db.commit()
                    elif e.error_code == 400:
                        # Invalid chat or user
                        pass
                    # Silently continue for other errors
                except Exception:
                    # Other errors - continue
                    pass
        
        # Send greetings grouped by chat
        for chat_id, users_data in chat_users.items():
            try:
                if len(users_data) == 1:
                    # Single user - use simple format
                    first_name, theme, user_id = users_data[0]
                    greeting = format_birthday_greeting(first_name, theme, user_id)
                    bot.send_message(chat_id, greeting, parse_mode='HTML')
                else:
                    # Multiple users - use special format
                    greeting = format_multiple_birthdays_greetings(users_data)
                    bot.send_message(chat_id, greeting, parse_mode='HTML')
            except telebot.apihelper.ApiTelegramException as e:
                # Handle errors: bot removed from chat, etc.
                if e.error_code == 403:
                    # Bot was blocked or removed from chat
                    chat = db.query(Chat).filter(Chat.chat_id == chat_id).first()
                    if chat:
                        chat.bot_is_admin = False
                        db.commit()
                # Silently continue for other errors
            except Exception:
                # Other errors - continue
                pass
    finally:
        db.close()


def setup_scheduler(bot: telebot.TeleBot) -> BlockingScheduler:
    """Setup and start the scheduler for daily birthday notifications."""
    # Parse notification time
    time_str: str = Config.NOTIFICATION_TIME
    try:
        hour, minute = map(int, time_str.split(':'))
    except (ValueError, AttributeError):
        hour, minute = 9, 0  # Default to 9:00
    
    # Get timezone
    timezone_str: str = Config.TIMEZONE
    try:
        tz = pytz.timezone(timezone_str)
    except (pytz.exceptions.UnknownTimeZoneError, AttributeError):
        tz = pytz.UTC
    
    # Create scheduler
    scheduler = BlockingScheduler(timezone=tz)
    
    # Add daily job
    scheduler.add_job(
        send_birthday_notifications,
        trigger=CronTrigger(hour=hour, minute=minute),
        args=[bot],
        id='daily_birthday_notifications',
        name='Send daily birthday notifications',
        replace_existing=True
    )
    
    # Add monthly job (1st day of each month)
    scheduler.add_job(
        send_monthly_birthday_overview,
        trigger=CronTrigger(day=1, hour=hour, minute=minute),
        args=[bot],
        id='monthly_birthday_overview',
        name='Send monthly birthday overview',
        replace_existing=True
    )
    
    return scheduler


def get_birthdays_in_range_for_chat(db: Session, chat_id: int, start_date: date, days: int) -> Dict[str, List[str]]:
    """Get birthdays in the specified date range for users in the chat."""
    birthdays = defaultdict(list)
    
    # Get all users in this chat
    users = db.query(User).join(UserChat).filter(
        UserChat.chat_id == chat_id
    ).distinct().all()
    
    end_date = start_date + timedelta(days=days)
    
    for user in users:
        # Create birthday date for current year
        try:
            birthday_this_year = datetime(start_date.year, user.birthday_month, user.birthday_day).date()
        except ValueError:
            # Invalid date (e.g., Feb 29 in non-leap year)
            continue
        
        # Check if birthday falls in range
        if start_date <= birthday_this_year < end_date:
            date_str = f"{user.birthday_day:02d}.{user.birthday_month:02d}"
            birthdays[date_str].append(user.first_name)
        else:
            # Check next year if we're near year end
            try:
                birthday_next_year = datetime(start_date.year + 1, user.birthday_month, user.birthday_day).date()
                if start_date <= birthday_next_year < end_date:
                    date_str = f"{user.birthday_day:02d}.{user.birthday_month:02d}"
                    birthdays[date_str].append(user.first_name)
            except ValueError:
                pass
    
    return birthdays


def send_monthly_birthday_overview(bot: telebot.TeleBot) -> None:
    """Send monthly birthday overview (like /month command) to all chats on 1st of each month."""
    db = get_db_session()
    try:
        today = datetime.now().date()
        
        # Get all chats where bot is admin
        admin_chats = db.query(Chat).filter(Chat.bot_is_admin == True).all()
        
        for chat in admin_chats:
            try:
                # Get birthdays for next 31 days
                birthdays = get_birthdays_in_range_for_chat(db, chat.chat_id, today, days=31)
                
                if not birthdays:
                    # Don't send message if no birthdays
                    continue
                
                # Format message (same as /month command)
                message_text = "🎂 Дни рождения на ближайшие 31 день:\n\n"
                for date_str, names in sorted(birthdays.items()):
                    names_list = ", ".join(names)
                    message_text += f"• {date_str}: {names_list}\n"
                
                bot.send_message(chat.chat_id, message_text)
            except telebot.apihelper.ApiTelegramException as e:
                # Handle errors: bot removed from chat, etc.
                if e.error_code == 403:
                    # Bot was blocked or removed from chat
                    chat.bot_is_admin = False
                    db.commit()
                # Silently continue for other errors
            except Exception:
                # Other errors - continue
                pass
    finally:
        db.close()