import discord
import os
from dotenv import load_dotenv
from datetime import datetime
from tokenizer import Tokenizer
from model import RubyTrainer

# Load environment
load_dotenv()
TOKEN = os.getenv("DISCORD_TOKEN")

if not TOKEN:
    raise RuntimeError("Bot token not found in .env")

# Setup intents
intents = discord.Intents.default()
intents.message_content = True
intents.dm_messages = True
intents = intents

class Ruby(discord.Client):
    def __init__(self):
        super().__init__(intents=intents)
        self.tokenizer = Tokenizer()
        self.trainer = RubyTrainer(self.tokenizer)
        self.log_path = os.path.join("logs", "messages.log")
        os.makedirs("logs", exist_ok=True)

    async def on_ready(self):
        print(f"[READY] Logged in as {self.user} (ID: {self.user.id})")

    async def on_message(self, message: discord.Message):
        if message.author.id == self.user.id:
            return

        self.log_message(message)
        self.trainer.train_on_tokens_from_text(message.content.strip())

        reply = self.trainer.generate_reply()
        if reply.strip():
            await message.channel.send(reply)
        else:
            print("[REPLY] Skipped (empty)")



    def log_message(self, message: discord.Message):
        timestamp = datetime.utcnow().isoformat()
        log_entry = f"{timestamp} | {message.author.name} | {message.content.strip()}\n"

        with open(self.log_path, "a", encoding="utf-8") as f:
            f.write(log_entry)

        print(f"[LOGGED] {log_entry.strip()}")

    def train_on_message(self, message: discord.Message):
        text = message.content.strip()
        self.trainer.train_on_tokens_from_text(text)
        token_tensor = torch.tensor(tokens, dtype=torch.long)
        loss = train_on_tokens(self.model, tokens, self.optimizer, self.criterion, device="cpu")
        print(f"[TRAIN] Tokens: {tokens} | Loss: {loss:.4f}")

# Run Ruby

client = Ruby()
client.run(TOKEN)