import asyncio
import datetime
from os.path import isfile
import pickle
import re
import sys
import traceback

import discord
from discord.ext import commands
import markovify
import os
import uuid
from random import randint

import botutils
from botutils import Utils
import botchecks

import music
from music import Music

owner_id = '126215805053435904'
bot_token='SEKRUT' ## THIS MUST BE REMOVED IF CODE IS SHARED
bot = commands.Bot(command_prefix='$', description="Yes, I am called Locutus!", pm_help=True, self_bot=False)
bot.training_data = []
bot.STATE_SIZE = 2

utilsCog = botutils.Utils(bot)
musicCog = music.Music(bot)
bot.add_cog([utilsCog,musicCog])

@bot.event
async def on_ready():
    print('Logged in as')
    print(bot.user.name)
    #print(bot.user.id)
    print('------')
    await weather_bot()


@bot.event
async def on_message(message):
    bot.training_data.append(message)
    await bot.process_commands(message)

@bot.command(pass_context=True)
async def load_markov(ctx):
    if isfile('./markovdata.pkl'):
        with open('./markovdata.pkl', mode='rb') as f:
            bot.training_data = pickle.load(f)
    if not 'predata' in sys.argv:
        bot.loop.create_task(load())
    else:
        print("Loaded!")
        await bot.say("Loaded Markov data!")

@bot.command(pass_context=True)
async def chanMarkov(ctx, channel:discord.Channel=None, num:int=1):
    text = ''
    if channel is None:
        channel = ctx.message.channel
    for m in bot.training_data:
        if m.channel.id == channel.id:
            text += m.content + '\n'
    markov = markovify.NewlineText(text, state_size=bot.STATE_SIZE)
    out = ''
    for i in range(num):
        try:
            out += markov.make_sentence(tries=1000) + '\n\n'
        except:
            pass
    out = fix_mentions(out)
    await bot.say(out)  # @UndefinedVariable

@bot.command(pass_context=True)
async def servMarkov(ctx, num:int=1):
    loading = await bot.say("Generating, this might take a while...")
    text = ''
    for m in bot.training_data:
        if m.server is not None:
            if m.server.id == ctx.message.server.id:
                text += m.content + '\n'
    markov = markovify.NewlineText(text, state_size=bot.STATE_SIZE)
    out = ''
    for i in range(num):
        try:
            out += markov.make_sentence(tries=1000) + '\n\n'
        except:
            pass
    out = fix_mentions(out)
    await bot.edit_message(loading, out)  # @UndefinedVariable

@bot.command(pass_context=True)
async def membMarkov(ctx, member:discord.Member, num:int=1):
    text = ''
    for m in bot.training_data:
        if m.server is not None:
            if m.author.id == member.id and m.server.id == ctx.message.server.id:
                text += m.content + '\n'
    markov = markovify.NewlineText(text, state_size=bot.STATE_SIZE)
    out = ''
    for i in range(num):
        try:
            out += markov.make_sentence(tries=1000) + '\n\n'
        except:
            pass
    out = fix_mentions(out)
    await bot.say(out)  # @UndefinedVariable

@bot.command(pass_context=True)
async def userMarkov(ctx, member:str, num:int=1, scrub="don't scrub"):
    text = ''
    for m in bot.training_data:
        if m.author.id == member:
            text += m.content + '\n'
    markov = markovify.NewlineText(text, state_size=bot.STATE_SIZE)
    out = ''
    for i in range(num):
        try:
            out += markov.make_sentence(tries=1000) + '\n\n'
        except:
            pass
    out = fix_mentions(out)
    if scrub == "scrub": out = scrub_names(out)
    await bot.say(out)  # @UndefinedVariable

@bot.command(pass_context=True)
async def globMarkov(ctx, num:int=1):
    print("Starting global markov...")
    text = ''
    for m in bot.training_data:
        text += m.content + '\n'
    markov = markovify.NewlineText(text, state_size=bot.STATE_SIZE)
    out = ''
    for i in range(num):
        try:
            out += markov.make_sentence(tries=1000) + '\n\n'
        except:
            pass
    out = fix_mentions(out)
    await bot.say(out)  # @UndefinedVariable

def fix_mentions(string):
    mentions = re.findall('<[@!]+[0-9]+>', string)
    for mention in mentions:
        string = string.replace(mention, "**@" + str(discord.utils.get(bot.get_all_members(), id=re.sub('[<>@!]', '', mention)))+"**")
    return string.replace('@', '@\u200b')

@bot.command(pass_context=True)
async def code(ctx, *, code : str):
    """Arbitrarily runs code."""
    if not ctx.message.author.id == '126215805053435904': return
    here = ctx.message.channel
    this = ctx.message
    def echo(out):
        msg(here, out)
    def rep(out):
        replace(this, out)
    def coro(coro):
        asyncio.ensure_future(coro)
    try:
        exec(code, globals(), locals())
    except:
        traceback.print_exc()
        out = discord_trim(traceback.format_exc())
        for o in out:
            await bot.send_message(ctx.message.channel, o)

def msg(dest, out):
    coro = bot.send_message(dest, out)
    asyncio.ensure_future(coro)

def replace(msg, out):
    coro = bot.edit_message(msg, out)
    asyncio.ensure_future(coro)

def discord_trim(str):
    result = []
    trimLen = 0
    lastLen = 0
    while trimLen <= len(str):
        trimLen += 1999
        result.append(str[lastLen:trimLen])
        lastLen += 1999
    return result

@bot.command()
async def state_size(size:int=2):
    bot.STATE_SIZE = size
    await bot.say("State size set to " + str(size))

@bot.command(pass_context=True)
async def load_all(ctx, chan:discord.Channel):
    if not ctx.message.author.id == '126215805053435904': return
    existing_data = [m.id for m in bot.training_data]
    await bot.say("Loading ALL data for channel " + str(chan))
    print("Loading ALL data for channel " + str(chan))
    if chan.type is discord.ChannelType.text:
        try:
            oldest_timestamp = chan.created_at
            all_loaded=0
            while True:
                loaded = 0
                loglen = 0
                async for m in bot.logs_from(chan, 10000, after=oldest_timestamp):
                    loglen += 1
                    if m.timestamp > oldest_timestamp:
                        oldest_timestamp = m.timestamp
                    if m.id not in existing_data:
                        bot.training_data.append(m)
                        loaded += 1
                        all_loaded += 1
                await asyncio.sleep(0.5)  # @UndefinedVariable
                print("... loaded {} messages.".format(str(loaded)))
                if loglen == 0: break
        except: pass
    print("Done! Loaded {} messages total.".format(str(all_loaded)))
    await bot.say("Done! Loaded {} messages total.".format(str(all_loaded)))

@bot.command(pass_context=True)
async def load_server(ctx):
    if not ctx.message.author.id == '126215805053435904': return
    existing_data = [m.id for m in bot.training_data]
    print("Loading ALL data for server " + str(ctx.message.server))
    all_loaded=0
    for chan in ctx.message.server.channels:
        print("Loading ALL data for channel " + str(chan))
        if chan.type is discord.ChannelType.text:
            try:
                oldest_timestamp = chan.created_at
                chan_loaded = 0
                while True:
                    loaded = 0
                    loglen = 0
                    async for m in bot.logs_from(chan, 10000, after=oldest_timestamp):
                        loglen += 1
                        if m.timestamp > oldest_timestamp:
                            oldest_timestamp = m.timestamp
                        if m.id not in existing_data:
                            bot.training_data.append(m)
                            loaded += 1
                            all_loaded += 1
                            chan_loaded += 1
                    await asyncio.sleep(0.5)  # @UndefinedVariable
                    print("... loaded {} messages.".format(str(loaded)))
                    if loglen == 0: break
            except: pass
        print("Done for channel {}! Loaded {} messages total.".format(str(chan), str(chan_loaded)))
        await bot.say("Done for channel {}! Loaded {} messages total.".format(str(chan), str(chan_loaded)))
    print("Done! Loaded {} messages total.".format(str(all_loaded)))
    await bot.say("Done! Loaded {} messages total.".format(str(all_loaded)))

async def load():
    await bot.wait_until_ready()
    existing_data = [m.id for m in bot.training_data]
    for s in bot.servers:
        print("Loading data for server " + str(s))
        for c in s.channels:
            print("Loading data for channel " + str(c))
            if c.type is discord.ChannelType.text:
                try:
                    loaded = 0
                    async for m in bot.logs_from(c, 1000):
                        if m.id not in existing_data:
                            bot.training_data.append(m)
                            loaded += 1
                except: pass
                print("... loaded {} messages.".format(str(loaded)))
            await asyncio.sleep(0.5)  # @UndefinedVariable
    print("Loaded!")

async def weather_bot():
    while True:
        try:
            await bot.change_presence(game=discord.Game(name=Utils.get_weather("short")))
        except Exception as e:
            print("Error: %s" % repr(e))
        await asyncio.sleep(1800)

try:
    bot.loop.run_until_complete(bot.start(bot_token, bot=True))
except KeyboardInterrupt:
    bot.loop.run_until_complete(bot.logout())
#     with open('./markovdata.pkl', mode='wb') as f:
#         pickle.dump(bot.training_data, f)
    temp = '%s-%s.tmp' % (uuid.uuid4(), 'markovdata.pkl')
    with open(temp, mode='wb') as tmp:
        pickle.dump(bot.training_data, tmp)

    # atomically move the file
    os.replace(temp, 'markovdata.pkl')
finally:
    bot.loop.close()