import os
import torch
import torch.nn.functional as F
import discord
from discord import Intents

from sensory_system.eyes import Eyes
from nervous_system.cortex import Cortex
from nervous_system.meta_learning import MetaLearner
from memory.hippocampus import Hippocampus
from motor_system.motor_cortex import MotorCortex
from headspace.dashboard import run_dashboard


class Organism:
    def __init__(self) -> None:
        # Device
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Sensory organ
        self.eyes = Eyes(books_path="content/books")

        # Memory & learning
        self.memory = Hippocampus()
        self.nervous_system = Cortex(...).to(self.device)
        self.meta = MetaLearner(self.nervous_system)
        self.motor = MotorCortex()

        # (Optional) Pre-load your 21+ books for future pre-training:
        self._load_corpus("content/books")

    def _load_corpus(self, folder_path: str) -> None:
        """Read all text files in content/books into memory for later use."""
        self.corpus = []
        for fn in os.listdir(folder_path):
            if fn.lower().endswith(".txt"):
                with open(os.path.join(folder_path, fn), encoding="utf-8") as f:
                    self.corpus.append(f.read())

    def learn_and_respond(self, message: str) -> str:
        # 1) Perception via eyes
        input_ids = self.eyes.preprocess(message)
        input_tensor = torch.tensor([input_ids], dtype=torch.long, device=self.device)

        # 2) Inference
        logits = self.nervous_system(input_tensor)
        response_ids = logits.argmax(dim=-1)[0].tolist()
        response = self.motor.decode(response_ids)

        # 3) Self-supervised loss (predict input back)
        loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)),
            input_tensor.view(-1),
        )

        # 4) Online meta-learning update
        self.meta.meta_update(loss)

        # 5) Store interaction
        self.memory.store({
            "input_ids": input_tensor.cpu(),
            "output_ids": response_ids,
            "input_text": message,
            "output_text": response
        })

        return response


# ————— Discord setup (all in one “body” file) —————
intents = Intents.default()
intents.message_content = True
client = discord.Client(intents=intents)
organism = Organism()


@client.event
async def on_ready() -> None:
    print(f"Logged in as {client.user}")


@client.event
async def on_message(message: discord.Message) -> None:
    if message.author == client.user or not message.content:
        return
    reply = organism.learn_and_respond(message.content)
    await message.channel.send(reply)


if __name__ == "__main__":
    TOKEN = os.getenv("DISCORD_TOKEN")
    if not TOKEN:
        raise RuntimeError("DISCORD_TOKEN environment variable not set.")
    run_dashboard(organism, host="0.0.0.0", port=5000)
    client.run(TOKEN)