import random

import discord
from discord.ext import commands

from app import cfg
from app.system import SystemPrompt
from app.var import CustomError
from pkg.llm import LanguageModel
from pkg.llm.data import Payload


class AurochsIntellect(commands.Cog):
    def __init__(self, bot: commands.Bot, llm_api: str):
        self.bot = bot
        self.llm = LanguageModel(llm_api, CustomError.LLM_EXCEPTION)
        
        self.prompts = [p.value.variable for p in SystemPrompt]
        self.weights = [p.value.weight for p in SystemPrompt]
        
    @commands.Cog.listener()
    async def on_message(self, message: discord.Message):
        assert self.bot.user is not None
        
        if message.author == self.bot.user or not self.bot.user.mentioned_in(message):
            return
            
        async with message.channel.typing():
            clean_content = message.content.replace(f'<@{self.bot.user.id}>', '').strip()

            system_prompt = random.choices(self.prompts, self.weights, k=1)[0]

            payload = Payload(
                system_prompt=system_prompt.format(message.author.display_name),
                user_message=clean_content,
                temperature=0.7,
                max_tokens=200,
                stream=False
            )
            
            response = await self.llm.make_request(payload)
            await message.reply(response)
            
    @commands.command()
    async def status(self, context: commands.Context):
        await context.send("The herd is accounted for.")
        
async def setup(bot: commands.Bot):
    await bot.add_cog(AurochsIntellect(bot, cfg.LLM_API))