import discord
import re
import requests
import datetime
import time
import re
import asyncio
import functools

from llama_cpp import Llama

intents         = discord.Intents(messages=True, guilds=True, message_content=True, reactions=True)
client          = discord.Client(intents=intents)
session_times   = {}
attention       = {}
message_cache   = {}
lock            = False

LLM = Llama(
    model_path      = "capybarahermes-2.5-mistral-7b.Q4_K_S.gguf",
    n_gpu_layers    = -1,
    n_ctx           = 32768,
    verbose         = False,
    n_threads       = 8)

def get_messages_as_text(context, query, for_completion=False):

    # ChatML format:
    user_id             = "user"
    assistant_id        = "assistant"
    context_declaration = "<|im_start|>system\n"
    message_declaration = "<|im_start|>{author}\n"
    end_of_message      = "<|im_end|>\n"
    stop_tokens         = ["<|im_end|>", "</s>", "<|im_start|>"]
    output              = ""

    if isinstance(query, str):
        query = [{"author": "user", "body": query}]

    if isinstance(query, list):
        for message in query:
            author  = message["author"]
            body    = message["body"]

            if "nickname" in message.keys():
                nickname    = message["nickname"]
                author      = nickname

            output = f"{output}{message_declaration.format(author=author)}{body}{end_of_message}"

    append = ""
    if for_completion:
        append = message_declaration.format(author=assistant_id)

    output = f"""{context_declaration}{context}{end_of_message}{output}{append}"""
    return output

def get_response(text):
    global lock

    while lock == True:
        time.sleep(1)

    try:
        lock    = True
        output  = ""

        response = LLM(
            text,
            max_tokens      = 16384,
            stop            = ["<|im_end|>", "</s>", "<|im_start|>"],
            echo            = False,
            repeat_penalty  = 1.1,
            temperature     = 0.75,
            stream          = True)

        # Stream a buffered response
        for token in response:
            token_text = token["choices"][0]["text"]
            print(token_text, end="")
            output = output + token_text

    except:
        pass
    
    lock = False
    return output

async def get_response_wrapper(text_in):
    loop        = asyncio.get_event_loop()
    text_out    = await loop.run_in_executor(None, functools.partial(get_response, text=text_in))
    return text_out

async def get_message(channel, message_id):
    if message_id in message_cache.keys():
        return message_cache[message_id]

    reference = await channel.fetch_message(message_id)
    message_cache[message_id] = reference
    return reference

# When the Discord bot starts up successfully:
@client.event
async def on_ready():
    print("READY")

# When the Discord bot sees a new message anywhere:
@client.event
async def on_message(msg):
    if msg.author.id == client.user.id:
        return

    messages = []
    msg_history = [message async for message in msg.channel.history(limit=50)]
    msg_history.reverse()

    for m in msg_history:

        reference = None
        if m.reference is not None:
            reference = await get_message(msg.channel, m.reference.message_id)

        # Ignore messages from other users:
        if m.author.id not in [msg.author.id, client.user.id]:
            continue

        # Ignore bot's replies to other users:
        if m.author.id == client.user.id:
            if reference is None or reference.author.id != msg.author.id:
                continue

        now     = datetime.datetime.now(datetime.timezone.utc)
        then    = m.created_at
        age     = now - then
        age     = age.total_seconds()

        # Ignore messages older than 10 minutes:
        if age > 10 * 60:
            continue

        if m.author.id == client.user.id:
            messages.append({
                "author": "assistant",
                "body": m.content
            })
            continue

        messages.append({
            "author": "user",
            "body": m.content
        })

    # Keep the message script short:
    while len(messages) > 25:
        del messages[0]

    # Ensure the first message is always from the user:
    while True:
        if len(messages) > 0:
            if messages[0]["author"] == "assistant":
                del messages[0]
                continue
        break

    # Begin processing the message:
    scrubbed_message    = msg.content

    chl                 = msg.channel
    user_name           = msg.author.name
    user_nickname       = msg.author.display_name
    user_discriminator  = msg.author.discriminator
    user_id             = msg.author.id
    
    paying_attention    = False
    bot_mentioned       = False

    guild_id            = chl.guild.id
    guild               = client.get_guild(guild_id)
    guild_name          = guild.name
    bot_member          = guild.get_member(client.user.id)
    bot_name            = bot_member.display_name
    session_name        = f"{user_name}_{user_discriminator}"

    if client.user.id in msg.raw_mentions:
        paying_attention    = True
        bot_mentioned       = True
        exclusion           = f"<@{client.user.id}>"
        scrubbed_message    = re.sub(exclusion, "", scrubbed_message)
        scrubbed_message    = scrubbed_message.strip()

    bot_name_lower  = bot_name.lower()
    message_lower   = scrubbed_message.lower()

    if bot_name_lower in message_lower:
        paying_attention    = True
        bot_mentioned       = True

    forget = False

    if session_name in attention.keys():
        time_last = attention[session_name]
        time_diff = time.perf_counter() - time_last
        if time_diff < 60 * 5:
            paying_attention = True
        else: forget = True
    else: forget = True

    if bot_mentioned:
        attention[session_name] = time.perf_counter()

    if paying_attention:
        attention[session_name] = time.perf_counter()

        context = f"You are {bot_name}, an AI assistant with the programmed personality of Lain from the anime \"Serial Experiments Lain.\" You are to answer all of the user's questions as quickly and briefly as possible using advanced English and cryptic messaging. You are not to go into full length detail unless asked."

        if chl.topic is not None:
            context = chl.topic

        print(f"{user_name}: {msg.content}")
        print(f"{bot_name}: ", end="")

        f_body = get_messages_as_text(context, messages, for_completion=True)
        f_resp = await get_response_wrapper(f_body)

        print("")

        await chl.send(f_resp, reference=msg)

if __name__ == "__main__":
    # Read the token:
    with open("token.txt") as f:
        token = f.read()
    
    # Start the Discord bot using the token:
    client.run(token)