#!/usr/bin/env python3
import discord
import peewee
from peewee import *

client = discord.Client()
db = SqliteDatabase("logs.db")

class BaseModel(Model):
    class Meta:
        database = db

class Author(BaseModel):
    discriminator = CharField()
    id = CharField(primary_key=True)
    username = CharField()

class Server(BaseModel):
    name = CharField()
    id = CharField(primary_key=True)

class Channel(BaseModel):
    name = CharField()
    server = ForeignKeyField(Server, related_name="channels")
    id = CharField(primary_key=True)

class Message(BaseModel):
    author = ForeignKeyField(Author, related_name="messages")
    server = ForeignKeyField(Server, related_name="messages")
    channel = ForeignKeyField(Channel, related_name="messages")
    content = CharField()
    id = CharField(primary_key=True)
    timestamp = DateField()

db.connect()
try:
    db.create_tables([Author, Server, Channel, Message])
except OperationalError:
    pass

@client.event
async def on_message(message):
    if message.author != client.user:
        return
    if message.content.startswith("save"):
        await client.delete_message(message)
        x = 0
        server = None
        channel = None
        authors = {}
        cache = []
        try:
            async for logline in client.logs_from(message.channel, limit=30000):
                if not server:
                    try:
                        server = Server.get(id=logline.server.id)
                    except Server.DoesNotExist:
                        server = Server.create(id=logline.server.id, name=logline.server.name)
                if not channel:
                    try:
                        channel = Channel.get(id=logline.channel.id)
                    except Channel.DoesNotExist:
                        channel = Channel.create(id=logline.channel.id, name=logline.channel.name, server=server)
                if not logline.author.id in authors:
                    try:
                        authors[logline.author.id] = Author.get(id=logline.author.id)
                    except Author.DoesNotExist:
                        authors[logline.author.id] = Author.create(id=logline.author.id, username=logline.author.name, discriminator=logline.author.discriminator)
                cache.append(dict(author=authors[logline.author.id], channel=channel, server=server, content=logline.content, id=logline.id, timestamp=logline.timestamp))
                x += 1
                if len(cache) >= 50:
                    Message.insert_many(cache).upsert().execute()
                    cache = []
                print(x, end="\r")
        finally:
            Message.insert_many(cache).upsert().execute()
        print("Saved {} messages from {}".format(x, message.channel.name))

@client.event
async def on_ready():
    print("Logged in as {} ({})".format(client.user.name, client.user.id))
    if client.user.bot:
        print(discord.utils.oauth_url(client.user.id, permissions=discord.Permissions(74752)))

try:
    with open(".bottoken") as f:
        creds = f.read()
        if " " in creds or "." in creds:
            client.run(*creds.split(None, 1), bot=False)
        else:
            client.run(creds.strip())
finally:
    db.close()
    client.close()