import discord
from discord.ext import commands
from discord import app_commands
import logging
import asyncio
import datetime
from typing import Optional, List, Dict, Any, Tuple
import asyncpg

# Configure logging
log = logging.getLogger(__name__)


class UserBannedError(commands.CheckFailure):
    """Custom exception for banned users."""

    def __init__(self, user_id: int, message: str):
        self.user_id = user_id
        self.message = message
        super().__init__(message)


class BanSystemCog(commands.Cog):
    """Cog for banning specific users from using the bot."""

    def __init__(self, bot: commands.Bot):
        self.bot = bot
        self.banned_users_cache = (
            {}
        )  # user_id -> {reason, message, banned_at, banned_by}

        # Create the main command group for this cog
        self.bansys_group = app_commands.Group(
            name="bansys", description="Bot user ban system commands (Owner only)"
        )

        # Register commands
        self.register_commands()

        # Add command group to the bot's tree
        self.bot.tree.add_command(self.bansys_group)

        log.info("BanSystemCog initialized with bansys command group.")

        # Setup database table when the cog is loaded
        self.bot.loop.create_task(self._setup_database())

        # Register the global check for prefix commands
        self.bot.add_check(self.check_if_user_banned)

        # Store the original interaction check if it exists
        self.original_interaction_check = getattr(
            self.bot.tree, "interaction_check", None
        )

        # Register our interaction check for slash commands
        self.bot.tree.interaction_check = self.interaction_check

    async def _setup_database(self):
        """Create the banned_users table if it doesn't exist."""
        # Wait for the bot to be ready to ensure the database pool is available
        await self.bot.wait_until_ready()

        if not hasattr(self.bot, "pg_pool") or self.bot.pg_pool is None:
            log.error(
                "PostgreSQL pool not available. Ban system will not work properly."
            )
            return

        try:
            async with self.bot.pg_pool.acquire() as conn:
                await conn.execute(
                    """
                    CREATE TABLE IF NOT EXISTS banned_users (
                        user_id BIGINT PRIMARY KEY,
                        reason TEXT,
                        message TEXT NOT NULL,
                        banned_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
                        banned_by BIGINT NOT NULL
                    );
                """
                )
                log.info("Created or verified banned_users table in PostgreSQL.")

                # Load banned users into cache
                await self._load_banned_users()
        except Exception as e:
            log.error(f"Error setting up banned_users table: {e}")

    async def _load_banned_users(self):
        """Load all banned users into the cache."""
        if not hasattr(self.bot, "pg_pool") or self.bot.pg_pool is None:
            log.error("PostgreSQL pool not available. Cannot load banned users.")
            return

        try:
            async with self.bot.pg_pool.acquire() as conn:
                records = await conn.fetch("SELECT * FROM banned_users")

                # Clear the current cache
                self.banned_users_cache.clear()

                # Populate the cache with the database records
                for record in records:
                    self.banned_users_cache[record["user_id"]] = {
                        "reason": record["reason"],
                        "message": record["message"],
                        "banned_at": record["banned_at"],
                        "banned_by": record["banned_by"],
                    }

                log.info(f"Loaded {len(records)} banned users into cache.")
        except Exception as e:
            log.error(f"Error loading banned users: {e}")

    @commands.Cog.listener()
    async def on_interaction(self, interaction: discord.Interaction):
        """Listener for all interactions to check if the user is banned."""
        # Skip check for the bot owner
        if interaction.user.id == self.bot.owner_id:
            return

        # Check if the user is banned
        if interaction.user.id in self.banned_users_cache:
            # Get the ban info
            ban_info = self.banned_users_cache[interaction.user.id]

            # If the interaction hasn't been responded to yet, respond with the ban message
            if not interaction.response.is_done():
                await interaction.response.send_message(
                    ban_info["message"], ephemeral=True
                )

            # Log the blocked interaction
            log.warning(
                f"Blocked interaction from banned user {interaction.user.id}: {ban_info['message']}"
            )

            # Raise the exception to prevent further processing
            raise UserBannedError(interaction.user.id, ban_info["message"])

    async def check_if_user_banned(self, ctx):
        """Global check to prevent banned users from using prefix commands."""
        # Skip check for DMs
        if not isinstance(ctx, commands.Context) and not hasattr(ctx, "guild"):
            return True

        # Get the user ID
        user_id = ctx.author.id if isinstance(ctx, commands.Context) else ctx.user.id

        # Check if the user is banned
        if user_id in self.banned_users_cache:
            ban_info = self.banned_users_cache[user_id]
            # Raise the custom exception with the ban message
            raise UserBannedError(user_id, ban_info["message"])

        # User is not banned, allow the command
        return True

    async def interaction_check(self, interaction: discord.Interaction) -> bool:
        """Global check for slash commands to prevent banned users from using them."""
        # Skip check for the bot owner
        if interaction.user.id == self.bot.owner_id:
            return True

        # Check if the user is banned
        if interaction.user.id in self.banned_users_cache:
            ban_info = self.banned_users_cache[interaction.user.id]

            # If the interaction hasn't been responded to yet, respond with the ban message
            if not interaction.response.is_done():
                try:
                    await interaction.response.send_message(
                        ban_info["message"], ephemeral=True
                    )
                except Exception as e:
                    log.error(
                        f"Error sending ban message to user {interaction.user.id}: {e}"
                    )

            # Raise the custom exception with the ban message
            raise UserBannedError(interaction.user.id, ban_info["message"])

        # If there was an original interaction check, call it
        if self.original_interaction_check is not None:
            return await self.original_interaction_check(interaction)

        # User is not banned, allow the interaction
        return True

    def register_commands(self):
        """Register all commands for this cog"""

        # --- Ban User Command ---
        ban_command = app_commands.Command(
            name="ban",
            description="Ban a user from using the bot",
            callback=self.bansys_ban_callback,
            parent=self.bansys_group,
        )
        app_commands.describe(
            user_id="The ID of the user to ban",
            message="The message to show when they try to use commands",
            reason="The reason for the ban (optional)",
            ephemeral="Whether the response should be ephemeral (only visible to the user)",
        )(ban_command)
        self.bansys_group.add_command(ban_command)

        # --- Unban User Command ---
        unban_command = app_commands.Command(
            name="unban",
            description="Unban a user from using the bot",
            callback=self.bansys_unban_callback,
            parent=self.bansys_group,
        )
        app_commands.describe(
            user_id="The ID of the user to unban",
            ephemeral="Whether the response should be ephemeral (only visible to the user)",
        )(unban_command)
        self.bansys_group.add_command(unban_command)

        # --- List Banned Users Command ---
        list_command = app_commands.Command(
            name="list",
            description="List all users banned from using the bot",
            callback=self.bansys_list_callback,
            parent=self.bansys_group,
        )
        app_commands.describe(
            ephemeral="Whether the response should be ephemeral (only visible to the user)"
        )(list_command)
        self.bansys_group.add_command(list_command)

    async def bansys_ban_callback(
        self,
        interaction: discord.Interaction,
        user_id: str,
        message: str,
        reason: Optional[str] = None,
        ephemeral: bool = True,
    ):
        """Ban a user from using the bot."""
        # Check if the user is the bot owner
        if interaction.user.id != self.bot.owner_id:
            await interaction.response.send_message(
                "This command can only be used by the bot owner.", ephemeral=ephemeral
            )
            return

        try:
            # Convert user_id to integer
            user_id_int = int(user_id)

            # Check if the user is already banned
            if user_id_int in self.banned_users_cache:
                await interaction.response.send_message(
                    f"User {user_id_int} is already banned.", ephemeral=ephemeral
                )
                return

            # Add the user to the database
            if hasattr(self.bot, "pg_pool") and self.bot.pg_pool is not None:
                async with self.bot.pg_pool.acquire() as conn:
                    await conn.execute(
                        """
                        INSERT INTO banned_users (user_id, reason, message, banned_by)
                        VALUES ($1, $2, $3, $4)
                        ON CONFLICT (user_id) DO UPDATE
                        SET reason = $2, message = $3, banned_by = $4, banned_at = CURRENT_TIMESTAMP
                    """,
                        user_id_int,
                        reason,
                        message,
                        interaction.user.id,
                    )

            # Add the user to the cache
            self.banned_users_cache[user_id_int] = {
                "reason": reason,
                "message": message,
                "banned_at": datetime.datetime.now(datetime.timezone.utc),
                "banned_by": interaction.user.id,
            }

            # Try to get the user's name for a more informative message
            try:
                user = await self.bot.fetch_user(user_id_int)
                user_display = f"{user.name} ({user_id_int})"
            except:
                user_display = f"User ID: {user_id_int}"

            await interaction.response.send_message(
                f"✅ Banned {user_display} from using the bot.\nMessage: {message}\nReason: {reason or 'No reason provided'}",
                ephemeral=ephemeral,
            )
            log.info(
                f"User {user_id_int} banned by {interaction.user.id}. Reason: {reason}"
            )

        except ValueError:
            await interaction.response.send_message(
                "Invalid user ID. Please provide a valid user ID.", ephemeral=ephemeral
            )
        except Exception as e:
            log.error(f"Error banning user {user_id}: {e}")
            await interaction.response.send_message(
                f"An error occurred while banning the user: {e}", ephemeral=ephemeral
            )

    async def bansys_unban_callback(
        self, interaction: discord.Interaction, user_id: str, ephemeral: bool = True
    ):
        """Unban a user from using the bot."""
        # Check if the user is the bot owner
        if interaction.user.id != self.bot.owner_id:
            await interaction.response.send_message(
                "This command can only be used by the bot owner.", ephemeral=ephemeral
            )
            return

        try:
            # Convert user_id to integer
            user_id_int = int(user_id)

            # Check if the user is banned
            if user_id_int not in self.banned_users_cache:
                await interaction.response.send_message(
                    f"User {user_id_int} is not banned.", ephemeral=ephemeral
                )
                return

            # Remove the user from the database
            if hasattr(self.bot, "pg_pool") and self.bot.pg_pool is not None:
                async with self.bot.pg_pool.acquire() as conn:
                    await conn.execute(
                        "DELETE FROM banned_users WHERE user_id = $1", user_id_int
                    )

            # Remove the user from the cache
            del self.banned_users_cache[user_id_int]

            # Try to get the user's name for a more informative message
            try:
                user = await self.bot.fetch_user(user_id_int)
                user_display = f"{user.name} ({user_id_int})"
            except:
                user_display = f"User ID: {user_id_int}"

            await interaction.response.send_message(
                f"✅ Unbanned {user_display} from using the bot.", ephemeral=ephemeral
            )
            log.info(f"User {user_id_int} unbanned by {interaction.user.id}.")

        except ValueError:
            await interaction.response.send_message(
                "Invalid user ID. Please provide a valid user ID.", ephemeral=ephemeral
            )
        except Exception as e:
            log.error(f"Error unbanning user {user_id}: {e}")
            await interaction.response.send_message(
                f"An error occurred while unbanning the user: {e}", ephemeral=ephemeral
            )

    async def bansys_list_callback(
        self, interaction: discord.Interaction, ephemeral: bool = True
    ):
        """List all users banned from using the bot."""
        # Check if the user is the bot owner
        if interaction.user.id != self.bot.owner_id:
            await interaction.response.send_message(
                "This command can only be used by the bot owner.", ephemeral=ephemeral
            )
            return

        if not self.banned_users_cache:
            await interaction.response.send_message(
                "No users are currently banned.", ephemeral=ephemeral
            )
            return

        # Create an embed to display the banned users
        embed = discord.Embed(
            title="Banned Users",
            description=f"Total banned users: {len(self.banned_users_cache)}",
            color=discord.Color.red(),
        )

        # Add each banned user to the embed
        for user_id, ban_info in self.banned_users_cache.items():
            # Try to get the user's name
            try:
                user = await self.bot.fetch_user(user_id)
                user_display = f"{user.name} ({user_id})"
            except:
                user_display = f"User ID: {user_id}"

            # Format the banned_at timestamp
            banned_at = (
                ban_info["banned_at"].strftime("%Y-%m-%d %H:%M:%S UTC")
                if isinstance(ban_info["banned_at"], datetime.datetime)
                else "Unknown"
            )

            # Try to get the banner's name
            try:
                banner = await self.bot.fetch_user(ban_info["banned_by"])
                banner_display = f"{banner.name} ({ban_info['banned_by']})"
            except:
                banner_display = f"User ID: {ban_info['banned_by']}"

            # Add a field for this user
            embed.add_field(
                name=user_display,
                value=f"**Reason:** {ban_info['reason'] or 'No reason provided'}\n"
                f"**Message:** {ban_info['message']}\n"
                f"**Banned at:** {banned_at}\n"
                f"**Banned by:** {banner_display}",
                inline=False,
            )

        await interaction.response.send_message(embed=embed, ephemeral=ephemeral)

    def cog_unload(self):
        """Cleanup when the cog is unloaded."""
        # Restore the original interaction check if it exists
        if hasattr(self, "original_interaction_check"):
            self.bot.tree.interaction_check = self.original_interaction_check
            log.info("Restored original interaction check on cog unload.")


# Setup function for loading the cog
async def setup(bot):
    """Add the BanSystemCog to the bot."""
    await bot.add_cog(BanSystemCog(bot))
    print("BanSystemCog setup complete.")