LichessStatTgWeb/LichessClientTG_bot/message_counters.py

"""
Module for tracking message counters in Telegram bot
"""
import sqlite3
import logging
from datetime import datetime, date
from typing import Dict, Any
from config import DATABASE_PATH

logger = logging.getLogger(__name__)


class MessageCounters:
    """Class for managing message counters"""
    
    def __init__(self, db_path: str = DATABASE_PATH):
        self.db_path = db_path
        self._ensure_counters_exist()
        self._reset_daily_counters_if_needed()
    
    def _ensure_counters_exist(self):
        """Ensure all command counters exist in database"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            commands = ['start', 'addgamer', 'addtoken', 'getgamers', 'delgamer', 
                       'today', 'yesterday', 'week', 'setperiod', 'lang', 'resetlang', 
                       'periodic_notification']
            for cmd in commands:
                cursor.execute('''
                    INSERT OR IGNORE INTO message_counters (command, total_count, today_count, last_reset_date)
                    VALUES (?, 0, 0, CURRENT_DATE)
                ''', (cmd,))
            conn.commit()
    
    def _reset_daily_counters_if_needed(self):
        """Reset daily counters if it's a new day"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            today = date.today()
            
            # Check all counters and reset if needed
            cursor.execute('''
                SELECT command, last_reset_date FROM message_counters
            ''')
            
            for row in cursor.fetchall():
                cmd, last_reset = row[0], row[1]
                if last_reset:
                    try:
                        last_reset_date = datetime.strptime(last_reset, '%Y-%m-%d').date()
                        if last_reset_date < today:
                            cursor.execute('''
                                UPDATE message_counters 
                                SET today_count = 0, last_reset_date = ?
                                WHERE command = ?
                            ''', (today.isoformat(), cmd))
                    except (ValueError, TypeError):
                        # Invalid date format, reset it
                        cursor.execute('''
                            UPDATE message_counters 
                            SET today_count = 0, last_reset_date = ?
                            WHERE command = ?
                        ''', (today.isoformat(), cmd))
                else:
                    # No reset date, set it to today
                    cursor.execute('''
                        UPDATE message_counters 
                        SET last_reset_date = ?
                        WHERE command = ?
                    ''', (today.isoformat(), cmd))
            
            conn.commit()
    
    def increment(self, command: str):
        """Increment counter for a command"""
        self._reset_daily_counters_if_needed()
        
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            
            # Ensure counter exists
            cursor.execute('''
                INSERT OR IGNORE INTO message_counters (command, total_count, today_count, last_reset_date)
                VALUES (?, 0, 0, CURRENT_DATE)
            ''', (command,))
            
            # Increment both counters
            cursor.execute('''
                UPDATE message_counters 
                SET total_count = total_count + 1,
                    today_count = today_count + 1
                WHERE command = ?
            ''', (command,))
            
            conn.commit()
    
    def get_all_stats(self) -> Dict[str, Any]:
        """Get all statistics"""
        self._reset_daily_counters_if_needed()
        
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT command, total_count, today_count
                FROM message_counters
                ORDER BY command
            ''')
            
            stats = {
                'by_command': {},
                'total_all_time': 0,
                'total_today': 0
            }
            
            for row in cursor.fetchall():
                cmd, total, today = row[0], row[1], row[2]
                stats['by_command'][cmd] = {
                    'total': total,
                    'today': today
                }
                stats['total_all_time'] += total
                stats['total_today'] += today
            
            return stats
    
    def get_stats_summary(self) -> Dict[str, Any]:
        """Get summary statistics for display"""
        stats = self.get_all_stats()
        
        return {
            'total_all_time': stats['total_all_time'],
            'total_today': stats['total_today'],
            'by_command': stats['by_command']
        }