import discord
import pymongo
from redbot.core import checks, commands
from redbot.core.bot import Red

class kBump(commands.Cog):
    """
    A bridge to the kBump MongoDB database.
    """

    def __init__(self, bot: Red):
        self.bot = bot
        self.client = None
        self.servers_collection = None
        # Create a background task to connect to the DB after the bot is ready
        self.db_init_task = self.bot.loop.create_task(self.initialize_db())

    async def initialize_db(self):
        """Initializes the database connection."""
        await self.bot.wait_until_ready()
        try:
            # Retrieve the secure URI we set with the '[p]set api' command
            uri = (await self.bot.get_shared_api_tokens("mongodb")).get("uri")
            if not uri:
                print("[kBump] MongoDB URI not set! Use '[p]set api mongodb uri ...' to set it.")
                return
            self.client = pymongo.MongoClient(uri)
            # Connect to the same database and collection as your other bot
            db = self.client["BytesBump"]
            self.servers_collection = db["servers"]
            print("[kBump] Successfully connected to the BytesBump MongoDB.")
        except pymongo.errors.PyMongoError as e:
            print(f"[kBump] PyMongo error connecting to BytesBump MongoDB: {e}")
            return
        except Exception as e:
            print(f"[kBump] General error connecting to BytesBump MongoDB: {e}")

def cog_unload(self):
    """Close the client connection when the cog is unloaded."""
    if self.client:
        try:
            self.client.close()
            print("[kBump] Closed connection to the BytesBump MongoDB.")
        except Exception as e:
            print(f"[kBump] Error closing connection to the BytesBump MongoDB: {e}")
    # Cancel the background task
    self.db_init_task.cancel()

@commands.command()
@commands.guild_only()
@checks.mod_or_permissions(manage_guild=True)
async def bumpstatus(self, ctx: commands.Context, guild: discord.Guild = None):
    """
    Checks the kBump setup status for a server.
    If no server is provided, it checks the current server.
    """
    try:
        if not self.servers_collection:
            await ctx.send("Database connection is not available. Please check the console logs.")
            return
        if not self.servers_collection:
            await ctx.send("The kBump database connection is not yet initialized. Please wait for the initialization task to complete.")
            return
        target_guild = guild or ctx.guild
        server_data = self.servers_collection.find_one({"_id": target_guild.id})
        if not server_data:
            await ctx.send(f"No kBump data found for the server: `{target_guild.name}`.")
            return
        # Rest of the code remains the same
    except pymongo.errors.PyMongoError as e:
        print(f"[kBump] PyMongo error getting kBump data for server: {e}")
        await ctx.send(f"Error getting kBump data for the server: `{target_guild.name}`. Please check the console logs for more information.")
        return
    except Exception as e:
        print(f"[kBump] General error getting kBump data for server: {e}")
        await ctx.send(f"Error getting kBump data for the server: `{target_guild.name}`. Please check the console logs for more information.")
        return

        embed_color = discord.Color(server_data.get('color', 0x2F3136))
        embed = discord.Embed(
                title=f"kBump Status for {target_guild.name}",
                color=embed_color
        )
        embed.set_thumbnail(url=target_guild.icon.url if target_guild.icon else None)

        description = server_data.get('description', 'Not set.')
        embed.add_field(name="📝 Description", value=f"```{description[:1000]}```", inline=False)
            
        custom_tags = server_data.get('tags', 'Not set.')
        embed.add_field(name="🏷️ Custom Tags", value=custom_tags, inline=False)

        discovery_tags = server_data.get('discovery_tags')
        if discovery_tags:
            embed.add_field(name="🌐 Official Discovery Tags", value=discovery_tags, inline=False)
            
        vanity = server_data.get('vanity_code')
        if vanity:
            embed.add_field(name="🔗 Vanity URL", value=f"`discord.gg/{vanity}`", inline=True)

        invite = server_data.get('invite_code')
        if invite:
            embed.add_field(name="✉️ Invite Code", value=f"`{invite}`", inline=True)

        await ctx.send(embed=embed)