import discord
import asyncio
import atexit
import os
import threading
from dotenv import load_dotenv
from datetime import datetime, timedelta
from dashboard import start_dashboard
from tokenizer import Tokenizer
from trainer import RubyTrainer
import logging

# Setup logging
logging.basicConfig(
    filename="logs/error.log",
    level=logging.ERROR,
    format="%(asctime)s %(levelname)s: %(message)s",
    encoding="utf-8"
)


# 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.last_message_time = datetime.utcnow()
        self.idle_threshold = timedelta(seconds=120)
        self.log_path = os.path.join("logs", "messages.log")
        os.makedirs("logs", exist_ok=True)

    async def setup_hook(self):
        self.loop.create_task(self.idle_dream_loop())

    async def set_activity(self, text=None):
        if text is None:
            await self.change_presence(status=discord.Status.online, activity=None)
        else:
            activity = discord.Activity(type=discord.ActivityType.listening, name=text)
            await self.change_presence(status=discord.Status.idle, activity=activity)

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

    async def idle_dream_loop(self):
        await self.wait_until_ready()
        while not self.is_closed():
            now = datetime.utcnow()
            if now - self.last_message_time > self.idle_threshold:
                print("[IDLE] Ruby has been idle — entering dream mode.")

                await self.set_activity("the past...")
                try:
                    self.trainer.dream()
                except Exception as e:
                    logging.error("Error dreaming: %s", e)

                await self.set_activity("my thoughts")
                from random import random
                speak = random() < 0.5
                thought = self.trainer.daydream(say_thought=speak)

                if speak and thought and len(thought.split()) >=4:
                    for guild in self.guilds:
                        for channel in guild.text_channels:
                            if channel.permissions_for(guild.me).send_messages:
                                if not thought.endswith("."):
                                    thought += "."
                                await channel.send(f"(dreaming) {thought}")
                                break
                        break  # only post to one server/channel

                await self.set_activity(None)  # reset to normal
                self.last_message_time = datetime.utcnow()

            await asyncio.sleep(180)

    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 = None
try:
    client = Ruby()

    def on_exit():
        if client:
            print("[EXIT] Ruby is gracefully shutting down...")
            client.trainer.dream()
            client.trainer.daydream(rounds=10)

    atexit.register(on_exit)
    dashboard_thread = threading.Thread(target=start_dashboard, daemon=True)
    dashboard_thread.start()
    client.run(TOKEN)
finally:
    if client is not None:
        print("[EXIT] Ruby is shutting down — dreaming one last time...")
        client.trainer.dream()
        client.trainer.daydream(rounds=10)