chatbot/__main__.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2020-2023 (c) Randy W @xtdevs, @xtsea
#
# from : https://github.com/TeamKillerX
# Channel : @RendyProjects
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import asyncio
import logging
import sys
import traceback
from contextlib import closing, suppress
from datetime import datetime as dt, timedelta

from pyrogram import idle
from uvloop import install

from chatbot import Randydev
from database import db
from logger import LOGS

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logging.getLogger("pyrogram.syncer").setLevel(logging.WARNING)
logging.getLogger("pyrogram.client").setLevel(logging.WARNING)
loop = asyncio.get_event_loop()

async def process_unfreezes(client):
    while True:
        try:
            now = dt.now()
            expired_users = db.user_blacklists.find({
                "is_frozen": True,
                "unfreeze_at": {"$lte": now}
            })
            
            async for user in expired_users:
                await db.user_blacklists.update_one(
                    {"_id": user["_id"]},
                    {"$set": {"is_frozen": False}}
                )
                LOGS.info(f"Auto-unfrozen user {user['user_id']}")

                with suppress(Exception):
                    await client.send_message(
                        user["user_id"],
                        "✅ Your restrictions have been automatically lifted"
                    )
            
            await asyncio.sleep(300)
            
        except Exception as e:
            LOGS.error(f"Unfreeze task error: {str(e)}")
            await asyncio.sleep(60)

async def shutdown(loop):
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
    for task in tasks:
        task.cancel()
    
    with suppress(asyncio.CancelledError):
        await asyncio.gather(*tasks, return_exceptions=True)
    
    LOGS.info("Application shutdown complete")

async def main():
    try:
        await db.connect()
        randydev = Randydev(loop=loop)
        
        await randydev.start()
        asyncio.create_task(process_unfreezes(randydev))
        LOGS.info("Application startup complete")

        await idle()

    except Exception as e:
        LOGS.critical(f"Fatal error: {str(e)}")
        traceback.print_exc()
    finally:
        await shutdown(asyncio.get_event_loop())

if __name__ == "__main__":
    install()

    with closing(loop):
        with suppress(asyncio.CancelledError, KeyboardInterrupt):
            loop.run_until_complete(main())

        loop.run_until_complete(asyncio.sleep(3))
        if not loop.is_closed():
            loop.close()