utils.py

"""Database utilities"""
import json
import csv
from datetime import datetime
from database import Database
import config


class DatabaseUtils:
    def __init__(self):
        self.db = Database()
    
    def export_chat_to_json(self, chat_id: int, output_file: str = None):
        """Export chat to JSON"""
        messages = self.db.get_messages(chat_id, limit=100000)
        
        if not output_file:
            output_file = f"export_chat_{chat_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(messages, f, ensure_ascii=False, indent=2, default=str)
        
        print(f"✓ Exported {len(messages)} messages to {output_file}")
        return output_file
    
    def export_chat_to_csv(self, chat_id: int, output_file: str = None):
        """Export chat to CSV"""
        messages = self.db.get_messages(chat_id, limit=100000)
        
        if not output_file:
            output_file = f"export_chat_{chat_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
        
        if not messages:
            print("⚠️ No messages to export")
            return None
        
        with open(output_file, 'w', encoding='utf-8', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            
            writer.writeheader()
            writer.writerows(messages)
        
        print(f"✓ Експортовано {len(messages)} повідомлень в {output_file}")
        return output_file
    
    def export_all_chats(self, format='json'):
        """Export all chats"""
        chats = self.db.get_chats(limit=10000)
        exported_files = []
        
        print(f"📦 Exporting {len(chats)} chats in {format} format...")
        
        for chat in chats:
            chat_id = chat['chat_id']
            print(f"  Exporting chat: {chat['title']}")
            
            if format == 'json':
                file = self.export_chat_to_json(chat_id)
            elif format == 'csv':
                file = self.export_chat_to_csv(chat_id)
            else:
                print(f"❌ Unknown format: {format}")
                continue
            
            if file:
                exported_files.append(file)
        
        print(f"\n✓ Total exported: {len(exported_files)} files")
        return exported_files
    
    def get_deleted_messages(self, chat_id: int = None):
        """Get all deleted messages"""
        with self.db.get_connection() as conn:
            cursor = conn.cursor()
            
            if chat_id:
                cursor.execute('''
                    SELECT m.*, c.title as chat_title
                    FROM messages m
                    JOIN chats c ON m.chat_id = c.chat_id
                    WHERE m.is_deleted = 1 AND m.chat_id = ?
                    ORDER BY m.deleted_at DESC
                ''', (chat_id,))
            else:
                cursor.execute('''
                    SELECT m.*, c.title as chat_title
                    FROM messages m
                    JOIN chats c ON m.chat_id = c.chat_id
                    WHERE m.is_deleted = 1
                    ORDER BY m.deleted_at DESC
                ''')
            
            return [dict(row) for row in cursor.fetchall()]
    
    def get_edited_messages(self, chat_id: int = None):
        """Get all edited messages"""
        with self.db.get_connection() as conn:
            cursor = conn.cursor()
            
            if chat_id:
                cursor.execute('''
                    SELECT m.*, c.title as chat_title
                    FROM messages m
                    JOIN chats c ON m.chat_id = c.chat_id
                    WHERE m.is_edited = 1 AND m.chat_id = ?
                    ORDER BY m.date DESC
                ''', (chat_id,))
            else:
                cursor.execute('''
                    SELECT m.*, c.title as chat_title
                    FROM messages m
                    JOIN chats c ON m.chat_id = c.chat_id
                    WHERE m.is_edited = 1
                    ORDER BY m.date DESC
                ''')
            
            return [dict(row) for row in cursor.fetchall()]
    
    def cleanup_old_media(self, days: int = 30):
        """Cleanup old media files (older than N days)"""
        import os
        from datetime import timedelta
        
        cutoff_date = datetime.now() - timedelta(days=days)
        deleted_count = 0
        freed_space = 0
        
        with self.db.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT media_path FROM messages
                WHERE media_path IS NOT NULL
                AND date < ?
            ''', (cutoff_date,))
            
            for row in cursor.fetchall():
                media_path = row[0]
                full_path = os.path.join(config.MEDIA_FOLDER, media_path)
                
                if os.path.exists(full_path):
                    size = os.path.getsize(full_path)
                    os.remove(full_path)
                    deleted_count += 1
                    freed_space += size
        
        freed_mb = freed_space / (1024 * 1024)
        print(f"✓ Deleted {deleted_count} files, freed {freed_mb:.2f} MB")
        return deleted_count, freed_space
    
    def get_chat_statistics(self, chat_id: int):
        """Detailed chat statistics"""
        with self.db.get_connection() as conn:
            cursor = conn.cursor()
            
            cursor.execute('SELECT COUNT(*) FROM messages WHERE chat_id = ?', (chat_id,))
            total = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(*) FROM messages WHERE chat_id = ? AND is_deleted = 1', (chat_id,))
            deleted = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(*) FROM messages WHERE chat_id = ? AND is_edited = 1', (chat_id,))
            edited = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(*) FROM messages WHERE chat_id = ? AND media_type IS NOT NULL', (chat_id,))
            with_media = cursor.fetchone()[0]
            
            cursor.execute('''
                SELECT sender_name, COUNT(*) as count
                FROM messages
                WHERE chat_id = ?
                GROUP BY sender_name
                ORDER BY count DESC
                LIMIT 10
            ''', (chat_id,))
            top_senders = [{'name': row[0], 'count': row[1]} for row in cursor.fetchall()]
            
            return {
                'total_messages': total,
                'deleted_messages': deleted,
                'edited_messages': edited,
                'messages_with_media': with_media,
                'top_senders': top_senders
            }


def main():
    """CLI для утиліт"""
    import sys
    
    utils = DatabaseUtils()
    
    if len(sys.argv) < 2:
        print("Usage:")
        print("  python utils.py export_json <chat_id>")
        print("  python utils.py export_csv <chat_id>")
        print("  python utils.py export_all [json|csv]")
        print("  python utils.py deleted [chat_id]")
        print("  python utils.py edited [chat_id]")
        print("  python utils.py cleanup <days>")
        print("  python utils.py stats <chat_id>")
        return
    
    command = sys.argv[1]
    
    if command == 'export_json':
        if len(sys.argv) < 3:
            print("❌ Specify chat_id")
            return
        chat_id = int(sys.argv[2])
        utils.export_chat_to_json(chat_id)
    
    elif command == 'export_csv':
        if len(sys.argv) < 3:
            print("❌ Specify chat_id")
            return
        chat_id = int(sys.argv[2])
        utils.export_chat_to_csv(chat_id)
    
    elif command == 'export_all':
        format = sys.argv[2] if len(sys.argv) > 2 else 'json'
        utils.export_all_chats(format)
    
    elif command == 'deleted':
        chat_id = int(sys.argv[2]) if len(sys.argv) > 2 else None
        messages = utils.get_deleted_messages(chat_id)
        print(f"\n📋 Found {len(messages)} deleted messages:\n")
        for msg in messages[:20]:
            print(f"[{msg['chat_title']}] {msg['sender_name']}: {msg['text'][:50]}")
    
    elif command == 'edited':
        chat_id = int(sys.argv[2]) if len(sys.argv) > 2 else None
        messages = utils.get_edited_messages(chat_id)
        print(f"\n📋 Found {len(messages)} edited messages:\n")
        for msg in messages[:20]:
            print(f"[{msg['chat_title']}] {msg['sender_name']}: {msg['text'][:50]}")
    
    elif command == 'cleanup':
        if len(sys.argv) < 3:
            print("❌ Specify number of days")
            return
        days = int(sys.argv[2])
        utils.cleanup_old_media(days)
    
    elif command == 'stats':
        if len(sys.argv) < 3:
            print("❌ Specify chat_id")
            return
        chat_id = int(sys.argv[2])
        stats = utils.get_chat_statistics(chat_id)
        
        print(f"\n📊 Chat {chat_id} statistics:\n")
        print(f"  Total messages: {stats['total_messages']}")
        print(f"  Deleted: {stats['deleted_messages']}")
        print(f"  Edited: {stats['edited_messages']}")
        print(f"  With media: {stats['messages_with_media']}")
        print(f"\n  Top senders:")
        for sender in stats['top_senders']:
            print(f"    {sender['name']}: {sender['count']} messages")
    
    else:
        print(f"❌ Unknown command: {command}")


if __name__ == '__main__':
    main()