from datetime import datetime, timezone


def shard(shard_id: int, shard_count: int) -> None:
    """Starts a new shard.

    Args:
        shard_id (int): The shard's ID.
        shard_count (int): The total amount of shards.
    """

    from payphone.bot import bot
    from payphone.core import config

    bot.shard_ids = [shard_id]
    bot.shard_count = shard_count

    bot.run(config.TOKEN)


if __name__ == "__main__":
    import multiprocessing as mp

    import time

    from tortoise import run_async

    import httpx

    from payphone.db import setup, models
    from payphone.core import config
    from payphone.core.redis import redis

    async def init():
        await setup.init()

        await models.Call.filter(ended_at=None).update(
            ended_at=datetime.now(timezone.utc)
        )
        await models.Channel.exclude(call_id=None).update(call_id=None)

        await redis.flushall()

    print("Running DB and Redis cleanup...")

    run_async(init())

    print("Cleanup done! Loading recommended shard amount...")

    response = httpx.get(
        "https://discord.com/api/v10/gateway/bot",
        headers={"Authorization": f"Bot {config.TOKEN}"},
    )

    info = response.json()

    shard_count = info["shards"]

    print(f"Starting {shard_count} shard{"s" if shard_count != 1 else ""}.")

    ps: list[mp.Process] = [
        mp.Process(target=shard, args=(i, shard_count)) for i in range(shard_count)
    ]

    for i, p in enumerate(ps):
        p.start()

        print(
            f"Shard {i}'s process started." + (" Waiting 5 seconds start initialize next shard..." if len(ps) - 1 < i else " All shards were started.")
        )

        if i + 1 != len(ps):
            time.sleep(5)

    try:
        while True:
            for i in range(shard_count):
                if not ps[i].is_alive():
                    print(f"Shard {i}'s process exited. Restarting the shard...")

                    ps[i] = mp.Process(target=shard, args=(i, shard_count))
                    ps[i].start()

                    print(
                        f"Shard {i} was restarted. Waiting 5 seconds to check next shard..."
                    )
                    time.sleep(5)

            time.sleep(5)

    except KeyboardInterrupt:
        print(f"Stopping {shard_count} shards...")

        for i, p in enumerate(ps):
            p.terminate()
            print(f"Terminated shard {i}.")

        print("Checking that all processes have finished...")
        
        for p in ps:
            p.join()

        print("All processes have been terminated.")