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

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

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_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)))
    target = client.get_channel(args.channel)
    if not target:
        print("Channel not found")
        await client.logout()
        return
    x = 0
    server = None
    channel = None
    authors = {}
    cache = []
    try:
        async for logline in client.logs_from(target, limit=args.limit):
            if logline.server:
                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, target.name))
    await client.logout()

try:
    parser = ArgumentParser()
    parser.add_argument("channel")
    parser.add_argument("--limit", default=10**15, type=int)
    args = parser.parse_args()
    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()