# Plugin for bot LLM chat
from discord.ext import commands
import discord
import io
import aiohttp
import yaml
import random
import os

plugin_folder=os.path.dirname(os.path.realpath(__file__))
prompts_folder=os.path.join(plugin_folder, 'prompts')
default_prompt="default.txt"
config_filename=os.path.join(plugin_folder, 'settings.yaml')
llm_data = {}

async def prompt_llm(prompt):
    print("Prompting LLM")
    print(f"PROMPT DATA\n{prompt}")
    async with aiohttp.ClientSession(llm_data["api_base"]) as session:
        async with session.post("/completion", json={"prompt": prompt, "n_predict": 250}) as resp:
            print(f"LLM response status {resp.status}")
            response_json=await resp.json()
            content=response_json["content"]
            return content

def get_message_contents(msg):
    message_text = f"{msg.author.name}: {msg.clean_content}"
    print(f"Message contents -- {message_text}")
    return message_text

async def get_chat_history(ctx, limit=20):
    messages = [message async for message in ctx.channel.history(limit=limit)]
    plain_messages = list(map(get_message_contents, messages))
    plain_messages.reverse()
    return plain_messages

@commands.command(name='llm')
async def llm_response(ctx):
    await ctx.channel.typing()
    prompt_file = os.path.join(prompts_folder, default_prompt)
    with open(prompt_file, 'r') as prompt_file:
        prompt = prompt_file.read()
    history_arr = await get_chat_history(ctx)
    history_str = '\n'.join(history_arr)
    full_prompt = prompt.replace("<CONVHISTORY>", history_str)
    response = await prompt_llm(full_prompt)
    await send_chat_responses(ctx, response)

async def send_chat_responses(ctx, response_text):
    print("Processing chat response")
    fullResponseLog = "dank-bot:" + response_text # first response won't include the user
    responseLines = fullResponseLog.splitlines()
    output_strs = []
    for line in responseLines:
        if line.startswith("dank-bot:"):
            truncStr = line.replace("dank-bot:","")
            output_strs.append(truncStr)
        elif line.find(":") > 0 and line.find(":") < 20:
            break
        else:
            output_strs.append(line.strip())
    for outs in output_strs:
        final_output_str = await fixup_mentions(ctx, outs)
        await ctx.channel.send(final_output_str)

async def fixup_mentions(ctx, text):
    newtext = text
    if (isinstance(ctx.channel,discord.DMChannel)):
        newtext = newtext.replace(f"@{ctx.author.name}", ctx.author.mention)
    elif (isinstance(ctx.channel,discord.GroupChannel)):
        for user in ctx.channel.recipients:
            newtext = newtext.replace(f"@{user.name}", user.mention)
    elif (isinstance(ctx.channel,discord.Thread)):
        for user in await ctx.channel.fetch_members():
            member_info = await ctx.channel.guild.fetch_member(user.id)
            newtext = newtext.replace(f"@{member_info.name}", member_info.mention)
    else:
        for user in ctx.channel.members:
            newtext = newtext.replace(f"@{user.name}", user.mention)
        if ctx.guild != None:
            for role in ctx.guild.roles:
                newtext = newtext.replace(f"@{role.name}", role.mention)
    return newtext

async def handle_message(ctx):
    print("Dank-bot received message")
    print(f"Dank-bot ID is {llm_data['bot'].user.id}")
    bot_id = llm_data['bot'].user.id

    # First case, bot DMed
    if (isinstance(ctx.channel,discord.DMChannel) and ctx.author.id != bot_id):
        print("Dank-bot DMed, responding")
        await llm_response(ctx)
        return

    # Second case, bot mentioned
    bot_mentions=list(filter(lambda x: x.id == bot_id, ctx.mentions))
    if (len(bot_mentions) > 0):
        print("Dank-bot mentioned, responding")
        await llm_response(ctx)
        return

    # Other case, random response
    random_roll = random.random()
    print(f"Dank-bot rolled {random_roll}")
    if (random_roll < llm_data['response_probability']):
        print(f"{random_roll} < {llm_data['response_probability']}, responding")
        await llm_response(ctx)
        return

async def setup(bot):
    with open(config_filename, 'r') as conf_file:
       yaml_config = yaml.safe_load(conf_file)
       llm_data["api_base"] = yaml_config["api_base"]
       llm_data["response_probability"] = yaml_config["response_probability"]
    bot.add_command(llm_response)
    bot.add_listener(handle_message, "on_message")
    llm_data["bot"] = bot
    print("LLM interface initialized")