import discord
from train_gpt_model import process_message
from gpt_model import load_model
import torch
from dotenv import load_dotenv
import os

load_dotenv()

# Discord bot token
TOKEN = os.getenv("DISCORD_TOKEN")

# Load the vocabulary
with open("vocab.txt", "r", encoding="utf-8") as f:
    text = f.read()
    chars = sorted(list(set(text)))

# Ensure that space and other special characters are included
required_chars = " \n\r\t"
for char in required_chars:
    if char not in chars:
        chars.append(char)

# Add a special token for unknown characters
special_token = "<unk>"
if special_token not in chars:
    chars.append(special_token)

vocab_size = len(chars)
string_to_int = {ch: i for i, ch in enumerate(chars)}
int_to_string = {i: ch for i, ch in enumerate(chars)}

# Initialize and load the model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = load_model(vocab_size, "phoebe_model.pt").to(device)

# Initialize Discord client
intents = discord.Intents.default()
intents.message_content = True
client = discord.Client(intents=intents)


@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

    # Debug: print the message content
    print(f"Received message: '{message.content}'")

    # Process the message and get a response
    response = process_message(message.content)

    # Send the response back to the Discord channel
    await message.channel.send(response)


client.run(TOKEN)