import os
import asyncio
import glob
import threading
import json
from collections import deque
import logging
import discord

from nervous_system import NervousSystem
import dashboard    # <-- import your new Flask app
import brain_map    # <-- import the blueprint to inject system

# Path for progress persistence
PROGRESS_PATH = 'progress.json'

# Mute logger
for noisy_logger in ["werkzeug", "flask", "flask.app"]:
    logging.getLogger(noisy_logger).setLevel(logging.CRITICAL)

# ─── Initialize Ruby & Discord ───────────────────────────────────────────────── # noqa: E501

intents = discord.Intents.default()
intents.message_content = True
client = discord.Client(intents=intents)
system = NervousSystem()
system.history = deque(maxlen=100)

# Load or resume vocab + embeddings
system.sensory.load_vocab('vocab.json')
system._resize_embeddings()
print('Loaded vocab size:', len(system.sensory.stoi))

# Resume progress
if os.path.isfile(PROGRESS_PATH):
    with open(PROGRESS_PATH, 'r', encoding='utf-8') as f:
        data = json.load(f)
        system.processed_lines = data.get('processed_lines', 0)
else:
    system.processed_lines = 0

# Compute total book lines
total = sum(
    1
    for path in glob.glob('books/*.txt')
    for line in open(path, encoding='utf-8')
    if line.strip()
)
system.total_lines = total

print(f'Resuming training at {system.processed_lines}/{system.total_lines} lines')

# Inject into Flask contexts
dashboard.system = system
brain_map.system = system

# ─── Book-training when idle ──────────────────────────────────────────────────── # noqa: E501


async def train_books_idle():
    await client.wait_until_ready()
    await asyncio.sleep(5)
    processed = 0
    skip = system.processed_lines

    for path in glob.glob('books/*.txt'):
        with open(path, encoding='utf-8') as f:
            for raw in f:
                text = raw.strip()
                if not text:
                    continue
                if processed < skip:
                    processed += 1
                    continue

                await asyncio.to_thread(system.train, text, text)
                processed += 1
                system.processed_lines = processed

                if processed % 200 == 0 or processed == system.total_lines:
                    system.sensory.save_vocab('vocab.json')
                    with open(PROGRESS_PATH, 'w', encoding='utf-8') as pf:
                        json.dump({'processed_lines': processed}, pf)

    # Final checkpoint
    system.sensory.save_vocab('vocab.json')
    with open(PROGRESS_PATH, 'w', encoding='utf-8') as pf:
        json.dump({'processed_lines': system.processed_lines}, pf)


@client.event
async def on_ready():
    print(f'Ruby is online as {client.user}!')
    asyncio.create_task(train_books_idle())


@client.event
async def on_message(message: discord.Message):
    if message.author == client.user or not message.content:
        return

    user_text = message.content.strip()
    reply = system.generate(user_text)
    await message.channel.send(reply)

    system.history.append({'user': user_text, 'bot': reply})
    asyncio.create_task(asyncio.to_thread(system.train, user_text, reply))


# ─── Launch Dashboard & Bot ──────────────────────────────────────────────────── # noqa: E501


def run_dashboard():
    dashboard.app.run(
        host='0.0.0.0', port=5000,
        debug=False, use_reloader=False
    )


threading.Thread(target=run_dashboard, daemon=True).start()
print('Dashboard available at http://127.0.0.1:5000')

token = os.getenv('DISCORD_TOKEN')
if not token:
    raise RuntimeError('Please set the DISCORD_TOKEN environment variable')
client.run(token)