import discord
import logging
import aiohttp
import io
from discord.ext import commands

logger = logging.getLogger(__name__)


class VoiceGeneratorProfile():
    def __init__(self, name, uuid, model="audio-cpp-chatterbox"):
        self.name = name
        self.model = model
        self.uuid = uuid


class VoiceCog(commands.Cog):
    def __init__(self, bot: commands.Bot):
        self._running = False
        self._current_voice = "Pricey"
        self.api_url = "http://10.6.9.4:8080"

    async def cog_load(self):
        data = await self._get_voices()
        self.voices = [VoiceGeneratorProfile(d['name'], d['id']) for d in data.get('data')]
        self._current_voice = self.voices[0]

    async def _get_voices(self):
        async with aiohttp.ClientSession() as session:
            data = await session.get(f"{self.api_url}/api/voice-profiles")
        data = await data.json()
        return data

    @commands.command(name="voice_set", description="Change the active voice by name")
    async def voice_set(self, ctx, voice_name):
        """Set the voice for the generator to use"""
        for v in self.voices:
            if v.name.lower() == voice_name.lower():
                self._current_voice = v
                await ctx.send(f"Swapped voice to {v.name}")
                return
        await ctx.send("Voice not found asshole")

    @commands.command(name="voice_list", description="List available voices")
    async def voice_list(self, ctx: commands.Context):
        """List available voices"""
        await ctx.channel.send(f"{"\n".join([f" - {v.name}" for v in self.voices])}")

    @commands.command(name="voice_info", description="See the status of the voice generator :tm:")
    async def voice_info(self, ctx: commands.Context):
        """View status of the voice generator"""
        data = f"""
 - Running: {self._running}
 - Current Voice: {self._current_voice.name}
"""
        await ctx.channel.send(data)

    @commands.command(name="voice_gen", description="Generate an audio clip using the currently configured voice")
    async def voice_gen(self, ctx: commands.Context, prompt: str):
        """Generate some voices"""
        if self._running:
            await ctx.send("All lines are currently busy, please try again later")
            return
        data = {"model": self._current_voice.model,
                "voice": f"localai://voice-profiles/{self._current_voice.uuid}",
                "input": prompt,
                "stream": False}
        await ctx.channel.send("Working on it!")
        self._running = True
        try:
            headers = {"Content-Type": "application/json"}
            async with aiohttp.ClientSession(headers=headers) as session:
                resp = await session.post(f"{self.api_url}/tts", json=data)
                d = await resp.read()
                # logging.debug(d)
                voice_data = io.BytesIO(d)
            self._running = False
            filedata = discord.File(voice_data, filename=f"{self._current_voice.name.replace(" ", "_").lower()}-{prompt[:10].replace(" ", "_").lower()}.wav")
            await ctx.channel.send(file=filedata)
        except Exception as e:
            self._running = False
            logging.error(e)


async def setup(bot: commands.Bot) -> None:
    """Register the voice_gen cog with the bot."""
    await bot.add_cog(VoiceCog(bot))