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

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

    def __init__(self, bot):
        self.bot = bot
        self.translator = Translator()
        self.languages = LANGUAGES

    @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
            
            # 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:
            # 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"

        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()}:
                await ctx.send(f"Invalid destination language code. Available languages: {', '.join(self.languages.keys())}")
                return

            # 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)

            # Create embed
            embed = discord.Embed(title="Translation", color=discord.Color.blue())
            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(embed=embed)

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

    @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)

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