#!/bin/python3.10
import sys
import os
import random

import discord
from discord.ext import commands
from llama_cpp import Llama
from dotenv import load_dotenv

description = """
An example bot to showcase the discord.ext.commands extension module.
There are a number of utility commands being showcased here.
"""

load_dotenv()
llm = Llama(model_path="./models/gpt4all-7B/gpt4all-lora-quantized.bin")

bot = commands.Bot(
    command_prefix=commands.when_mentioned_or("!"),
    description=description,
    intents=discord.Intents.all(),
)


@bot.event
async def on_ready():
    print(f"Logged in as {bot.user} (ID: {bot.user.id})")


@bot.event
async def on_message(message: discord.Message):
    msg = ""
    # Make sure we won't be replying to ourselves.
    if message.author.id == bot.user.id:
        return

    if f"""<@{bot.user.id}>""" in message.content:
        async with message.channel.typing():
            question = f"""Text transcript of a never ending dialog, where {message.author} interacts with Karl Marx.
        Karl is helpful, kind, honest, friendly, good at writing and never fails to answer {message.author}’s requests immediately and with details and precision.
        There are no annotations like (30 seconds passed...) or (to himself), just what {message.author} and Karl say aloud to each other.
        The dialog lasts for years, the entirety of it is shared below. It's 10000 pages long.

        {message.author}: What is Communism?
        Karl Marx: Communism is the doctrine of the conditions of the liberation of the proletariat.
        {message.author}: {message.content}"""
            msg = llm(question, max_tokens=256, stop=[f"""{message.author}"""], echo=True)[
                "choices"][0]["text"][len(question):]

        await message.channel.send(msg)

bot.run(os.environ.get("DISCORD_TOKEN"))