import discord
from discord import app_commands
from discord.ext import commands
import aiohttp
import random
import string
import os

from src.utils.db import get_db
from models import Message
from config import BOT_COLOR, LOG


class Archive(commands.Cog):
    def __init__(self, bot):
        self.bot = bot

    async def download_attachments(self, attachments) -> list:
        """Download attachments and return a list of their paths."""
        paths = []

        for attachment in attachments:
            async with aiohttp.ClientSession() as session:
                async with session.get(attachment.url) as response:
                    # Check if the request was successful
                    if response.status != 200:
                        LOG.warn(
                            f"Failed to download attachment: {attachment.url}"
                        )
                        continue

                    # Check for content type
                    content_type = response.headers.get("Content-Type")
                    if not content_type:
                        LOG.warn(
                            f"Failed to get content type for: {attachment.url}"
                        )
                        continue

                    # Create a randomized filename
                    file_extension = content_type.split("/")[-1]
                    filename = (
                        "".join(
                            random.choice(string.ascii_letters)
                            for i in range(10)
                        )
                        + f".{file_extension}"
                    )

                    # Save the attachment
                    with open(f"images/{filename}", "wb") as file:
                        file.write(await response.read())

                    # Add the path to the attachments list
                    paths.append(f"images/{filename}")

        return paths

    @app_commands.command()
    async def archive(
        self,
        interaction: discord.Interaction,
        channel: discord.TextChannel,
        amount: int,
    ):
        """Archive a channel's messages."""
        # Ensure valid channel permissions
        if not channel.permissions_for(
            interaction.guild.me
        ).read_message_history:
            return await interaction.response.send_message(
                "I do not have permission to read message history in that"
                " channel.",
                ephemeral=True,
            )

        # Ensure valid amount
        if amount < 1:
            return await interaction.response.send_message(
                "You must provide a number greater than 0.",
                ephemeral=True,
            )

        embed = discord.Embed(
            title="Archive Beginning",
            description=(
                f"Archiving {amount} messages from {channel.mention}. I will"
                " DM you a message once complete, make sure to allow messages"
                " from me."
            ),
            color=BOT_COLOR,
        )
        await interaction.response.send_message(embed=embed, ephemeral=True)

        db = next(get_db())
        count = 0
        messages = channel.history(limit=amount)
        async for message in messages:
            count += 1

            if not os.path.exists("images"):
                os.makedirs("images")

            attachments = await self.download_attachments(message.attachments)

            db_message = Message(
                timestamp=message.created_at.isoformat(),
                message_id=message.id,
                author_id=message.author.id,
                channel_id=channel.id,
                stickers=[sticker.name for sticker in message.stickers],
                role_mentions=[role.id for role in message.role_mentions],
                mention_everyone=message.mention_everyone,
                mentions=[mention.id for mention in message.mentions],
                attachments=attachments,
                content=message.content,
            )

            db.add(db_message)

            if count > 500:
                db.commit()
                count = 0

        db.commit()

        embed = discord.Embed(
            title="Archive Complete",
            description=f"Archived {amount} messages from {channel.mention}.",
            color=BOT_COLOR,
        )

        try:
            await interaction.user.send(embed=embed)
        except discord.Forbidden:
            await channel.send(
                f"{interaction.user.mention} I have completed the archive, but"
                " was unable to DM you the final message. Please check your"
                " DM settings to receive future messages."
            )


async def setup(bot):
    await bot.add_cog(Archive(bot))