from redbot.core import commands, Config
from googletrans import Translator, LANGUAGES
import discord
from discord import app_commands
import asyncio

class TranslatorCog(commands.Cog):
    """Translate messages using Google Translate"""

    def __init__(self, bot):
        self.bot = bot
        self.translator = Translator()
        self.languages = LANGUAGES
        self.config = Config.get_conf(self, identifier=95932766180)
        default_guild = {
            "use_embeds": True
        }
        self.config.register_guild(**default_guild)

    @commands.group(name="transset")
    @commands.admin_or_permissions(manage_guild=True)
    async def _transset(self, ctx):
        """Configure translator settings"""
        pass

    @_transset.command(name="embed")
    async def _transset_embed(self, ctx, toggle: bool):
        """Toggle whether to use embeds for translations or plain text"""
        await self.config.guild(ctx.guild).use_embeds.set(toggle)
        await ctx.send(f"Embed display has been {'enabled' if toggle else 'disabled'}.")

    async def translate_and_respond(self, ctx, text_to_translate, dest_lang="en", source_lang="auto", original_author=None, ephemeral=False):
        try:
            # Validate destination language
            if dest_lang.lower() not in self.languages and dest_lang.lower() not in {v.lower(): k for k, v in self.languages.items()}:
                return await ctx.send(f"Invalid destination language code. Use `{ctx.prefix}langlist` to see available languages.", ephemeral=True)

            # Convert language name to code if full name was provided
            if dest_lang.lower() in {v.lower(): k for k, v in self.languages.items()}:
                dest_lang = next(k for k, v in self.languages.items() if v.lower() == dest_lang.lower())

            # Translate the text
            translation = self.translator.translate(text_to_translate, dest=dest_lang, src=source_lang)
            
            use_embeds = await self.config.guild(ctx.guild).use_embeds()
            
            if use_embeds:
                embed = discord.Embed(title="Translation", color=discord.Color.blue())
                if original_author:
                    embed.set_author(name=f"Original message by {original_author.display_name}", icon_url=original_author.display_avatar.url)
                embed.add_field(name=f"Original ({translation.src}):", value=text_to_translate[:1024], inline=False)
                embed.add_field(name=f"Translation ({translation.dest}):", value=translation.text[:1024], inline=False)
                await ctx.send(content=ctx.author.mention, embed=embed, ephemeral=ephemeral)
            else:
                author_text = f"\nOriginal message by: {original_author.display_name}" if original_author else ""
                response = (
                    f"{ctx.author.mention}\n"
                    f"**Original** ({translation.src}):{author_text}\n{text_to_translate}\n\n"
                    f"**Translation** ({translation.dest}):\n{translation.text}"
                )
                await ctx.send(response, ephemeral=ephemeral)

        except Exception as e:
            await ctx.send(f"An error occurred while translating: {str(e)}", ephemeral=True)

    @commands.command(name="translate")
    async def translate_text(self, ctx, *, message=None):
        """Translate a message to English by replying to it
        You can also specify a language code after your message to translate to that language
        Example: [p]translate bonjour fr -> en
        """
        if not message and not ctx.message.reference:
            await ctx.send("Please either provide text to translate or reply to a message to translate it.")
            return

        # If replying to a message
        if ctx.message.reference:
            referenced_msg = await ctx.fetch_message(ctx.message.reference.message_id)
            text_to_translate = referenced_msg.content
            original_author = referenced_msg.author
            
            # Check if language was specified in the command
            if message:
                try:
                    source_lang, dest_lang = message.split("->")
                    source_lang = source_lang.strip()
                    dest_lang = dest_lang.strip()
                except ValueError:
                    dest_lang = "en"
                    source_lang = "auto"
            else:
                dest_lang = "en"
                source_lang = "auto"
        else:
            original_author = ctx.author
            # Direct message translation
            if "->" in message:
                try:
                    text, langs = message.rsplit("->", 1)
                    text_to_translate = text.strip()
                    dest_lang = langs.strip()
                    source_lang = "auto"
                except ValueError:
                    text_to_translate = message
                    dest_lang = "en"
                    source_lang = "auto"
            else:
                text_to_translate = message
                dest_lang = "en"
                source_lang = "auto"

        await self.translate_and_respond(ctx, text_to_translate, dest_lang, source_lang, original_author)

    @commands.command(name="langlist")
    async def language_list(self, ctx):
        """List all available language codes"""
        lang_list = []
        for code, lang in self.languages.items():
            lang_list.append(f"`{code}`: {lang}")
        
        # Split into chunks of 20 for readability
        chunks = [lang_list[i:i + 20] for i in range(0, len(lang_list), 20)]
        
        for i, chunk in enumerate(chunks):
            if i == 0:
                embed = discord.Embed(title="Available Languages", 
                                    description="Use these codes with the translate command\n\n" + "\n".join(chunk),
                                    color=discord.Color.blue())
            else:
                embed = discord.Embed(description="\n".join(chunk), color=discord.Color.blue())
            await ctx.send(embed=embed)

    @app_commands.command(name="translate")
    @app_commands.describe(
        text="The text to translate",
        to_language="The language to translate to (e.g., 'en' for English)",
        from_language="The language to translate from (optional)"
    )
    async def slash_translate(self, interaction: discord.Interaction, text: str, to_language: str, from_language: str = "auto"):
        """Translate text from one language to another"""
        await interaction.response.defer(ephemeral=True)
        await self.translate_and_respond(interaction, text, to_language, from_language, interaction.user, ephemeral=True)

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