# main.py: Discord Bot Code

import discord
import torch
from model import JadeModel
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

intents = discord.Intents.default()
intents.messages = True
intents.message_content = True
client = discord.Client(intents=intents)

# Initialize the model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = JadeModel().to(device)


@client.event
async def on_ready():
    print(f'We have logged in as {client.user}')


@client.event
async def on_message(message):
    if message.author == client.user:
        return

    # Train Jade with the new message
    model.train_on_message(message.content)

    # Generate a response using Jade
    response = model.generate_response(message.content)
    await message.channel.send(response)

# Start the bot with your token
client.run(os.getenv('DISCORD_TOKEN'))