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

import discord
from ollama import AsyncClient
from discord.ext import commands
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()
client = AsyncClient()

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})")

def extract_response(input_string):
    # Find the position of the closing tag
    end_index = input_string.find('</think>')

    if end_index == -1:
        raise ValueError("Closing tag </think> not found")

    # Extract the response part of the string
    response = input_string[end_index + len('</think>'):].strip()

    return response

@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():
            messages = [
                {
                    'role':'user',
                    'content': message.content 
                }
            ]

            response = await client.chat('marx', messages=messages)
            print(response['message']['content'])
            msg = extract_response(response['message']['content'])

        await message.channel.send(msg)

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