# functions.py | function definitions
# Copyright (C) 2019-2021  EraserBird, person_v1.32, hmmm

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import datetime
import itertools
import os
import pickle
import shutil
import string
from typing import Union, Iterable

import discord
from discord.ext import commands

import sciolyid.config as config
from sciolyid.data import (
    GenericError,
    all_categories,
    database,
    dealias_group,
    get_category,
    groups,
    id_list,
    logger,
    states,
)
from sciolyid.util import fetch_get_user


async def channel_setup(ctx):
    """Sets up a new discord channel.

    `ctx` - Discord context object
    """
    logger.info("checking channel setup")
    if not database.exists(f"channel:{ctx.channel.id}"):
        database.hset(
            f"channel:{ctx.channel.id}",
            mapping={"item": "", "answered": 1, "prevJ": 20, "prevI": ""},
        )
        # true = 1, false = 0, prevJ is 20 to define as integer
        logger.info("channel data added")
        await ctx.send("Ok, setup! I'm all ready to use!")

    if database.zscore("score:global", str(ctx.channel.id)) is None:
        database.zadd("score:global", {str(ctx.channel.id): 0})
        logger.info("channel score added")

    if ctx.guild is not None:
        database.zadd("channels:global", {f"{ctx.guild.id}:{ctx.channel.id}": 0})


async def user_setup(ctx):
    """Sets up a new discord user for score tracking.

    `ctx` - Discord context object
    """
    logger.info("checking user data")
    if database.zscore("users:global", str(ctx.author.id)) is None:
        database.zadd("users:global", {str(ctx.author.id): 0})
        logger.info("user global added")
        await ctx.send("Welcome <@" + str(ctx.author.id) + ">!")

    date = str(datetime.datetime.now(datetime.timezone.utc).date())
    if database.zscore(f"daily.score:{date}", str(ctx.author.id)) is None:
        database.zadd(f"daily.score:{date}", {str(ctx.author.id): 0})
        logger.info("user daily added")

    # Add streak
    if (database.zscore("streak:global", str(ctx.author.id)) is None) or (
        database.zscore("streak.max:global", str(ctx.author.id)) is None
    ):
        database.zadd("streak:global", {str(ctx.author.id): 0})
        database.zadd("streak.max:global", {str(ctx.author.id): 0})
        logger.info("added streak")

    if ctx.guild is not None:
        global_score = database.zscore("users:global", str(ctx.author.id))
        database.zadd(
            f"users.server:{ctx.guild.id}", {str(ctx.author.id): global_score}
        )
        logger.info("synced scores")


def item_setup(ctx, item: str):
    """Sets up a new item for incorrect tracking.

    `ctx` - Discord context object
    `item` - item to setup
    """
    logger.info("checking item data")
    if database.zscore("incorrect:global", string.capwords(item)) is not None:
        logger.info("item global ok")
    else:
        database.zadd("incorrect:global", {string.capwords(item): 0})
        logger.info("item global added")

    if (
        database.zscore(f"incorrect.user:{ctx.author.id}", string.capwords(item))
        is not None
    ):
        logger.info("incorrect item user ok")
    else:
        database.zadd(f"incorrect.user:{ctx.author.id}", {string.capwords(item): 0})
        logger.info("incorrect item user added")

    if (
        database.zscore(f"correct.user:{ctx.author.id}", string.capwords(item))
        is not None
    ):
        logger.info("correct item user ok")
    else:
        database.zadd(f"correct.user:{ctx.author.id}", {string.capwords(item): 0})
        logger.info("correct item user added")

    date = str(datetime.datetime.now(datetime.timezone.utc).date())
    if database.zscore(f"daily.incorrect:{date}", string.capwords(item)) is not None:
        logger.info("item daily ok")
    else:
        database.zadd(f"daily.incorrect:{date}", {string.capwords(item): 0})
        logger.info("item daily added")

    if database.zscore("frequency.item:global", string.capwords(item)) is not None:
        logger.info("item freq global ok")
    else:
        database.zadd("frequency.item:global", {string.capwords(item): 0})
        logger.info("item freq global added")

    if ctx.guild is not None:
        logger.info("no dm")
        if (
            database.zscore(f"incorrect.server:{ctx.guild.id}", string.capwords(item))
            is not None
        ):
            logger.info("item server ok")
        else:
            database.zadd(
                f"incorrect.server:{ctx.guild.id}", {string.capwords(item): 0}
            )
            logger.info("item server added")
    else:
        logger.info("dm context")

    if database.exists(f"session.data:{ctx.author.id}"):
        logger.info("session in session")
        if (
            database.zscore(f"session.incorrect:{ctx.author.id}", string.capwords(item))
            is not None
        ):
            logger.info("item session ok")
        else:
            database.zadd(
                f"session.incorrect:{ctx.author.id}", {string.capwords(item): 0}
            )
            logger.info("item session added")
    else:
        logger.info("no session")


def session_increment(ctx, item: str, amount: int = 1):
    """Increments the value of a database hash field by `amount`.

    `ctx` - Discord context object\n
    `item` - hash field to increment (see data.py for details,
    possible values include correct, incorrect, total)\n
    `amount` (int) - amount to increment by, usually 1
    """
    if database.exists(f"session.data:{ctx.author.id}"):
        logger.info("session active")
        logger.info(f"incrementing {item} by {amount}")
        value = int(database.hget(f"session.data:{ctx.author.id}", item))
        value += int(amount)
        database.hset(f"session.data:{ctx.author.id}", item, str(value))
    else:
        logger.info("session not active")


def incorrect_increment(ctx, item: str, amount: int = 1):
    """Increments the value of an incorrect item by `amount`.

    `ctx` - Discord context object\n
    `item` - item that was incorrect\n
    `amount` (int) - amount to increment by, usually 1
    """
    logger.info(f"incrementing incorrect {item} by {amount}")
    date = str(datetime.datetime.now(datetime.timezone.utc).date())
    database.zincrby("incorrect:global", amount, string.capwords(item))
    database.zincrby(f"incorrect.user:{ctx.author.id}", amount, string.capwords(item))
    database.zincrby(f"daily.incorrect:{date}", amount, string.capwords(item))
    if ctx.guild is not None:
        logger.info("no dm")
        database.zincrby(
            f"incorrect.server:{ctx.guild.id}", amount, string.capwords(item)
        )
    else:
        logger.info("dm context")
    if database.exists(f"session.data:{ctx.author.id}"):
        logger.info("session in session")
        database.zincrby(
            f"session.incorrect:{ctx.author.id}", amount, string.capwords(item)
        )
    else:
        logger.info("no session")


def score_increment(ctx, amount: int = 1):
    """Increments the score of a user by `amount`.

    `ctx` - Discord context object\n
    `amount` (int) - amount to increment by, usually 1
    """
    logger.info(f"incrementing score by {amount}")
    date = str(datetime.datetime.now(datetime.timezone.utc).date())
    database.zincrby("score:global", amount, str(ctx.channel.id))
    database.zincrby("users:global", amount, str(ctx.author.id))
    database.zincrby(f"daily.score:{date}", amount, str(ctx.author.id))
    if ctx.guild is not None:
        logger.info("no dm")
        database.zincrby(f"users.server:{ctx.guild.id}", amount, str(ctx.author.id))
    else:
        logger.info("dm context")
    if database.exists(f"race.data:{ctx.channel.id}"):
        logger.info("race in session")
        database.zincrby(f"race.scores:{ctx.channel.id}", amount, str(ctx.author.id))


def streak_increment(ctx, amount: int):
    """Increments the streak of a user by `amount`.

    `ctx` - Discord context object\n
    `amount` (int) - amount to increment by, usually 1.
    If amount is None, the streak is ended.
    """

    if amount is not None:
        # increment streak and update max
        database.zincrby("streak:global", amount, ctx.author.id)
        if database.zscore("streak:global", ctx.author.id) > database.zscore(
            "streak.max:global", ctx.author.id
        ):
            database.zadd(
                "streak.max:global",
                {ctx.author.id: database.zscore("streak:global", ctx.author.id)},
            )
    else:
        database.zadd("streak:global", {ctx.author.id: 0})


def check_state_role(ctx) -> list:
    """Returns a list of state roles a user has.

    `ctx` - Discord context object
    """
    logger.info("checking roles")
    if not config.options["state_roles"]:
        return []
    user_states = []
    if ctx.guild is not None:
        logger.info("server context")
        user_role_names = [role.name.lower() for role in ctx.author.roles]
        for state in states:
            # gets similarities
            if set(user_role_names).intersection(set(states[state]["aliases"])):
                user_states.append(state)
    else:
        logger.info("dm context")
    logger.info(f"user roles: {user_states}")
    return user_states


async def send_leaderboard(
    ctx, title, page, database_key=None, data=None, items_per_page=10
):
    logger.info("building/sending leaderboard")

    if database_key is None and data is None:
        raise GenericError("database_key and data are both NoneType", 990)
    if database_key is not None and data is not None:
        raise GenericError("database_key and data are both set", 990)

    if page < 1:
        page = 1

    entry_count = (
        int(database.zcard(database_key)) if database_key is not None else data.count()
    )
    page = (page * 10) - 10

    if entry_count == 0:
        logger.info(f"no items in {database_key}")
        await ctx.send("There are no items in the database.")
        return

    if page > entry_count:
        page = entry_count - (entry_count % 10)

    leaderboard_list = (
        map(
            lambda x: (x[0].decode("utf-8"), x[1]),
            database.zrevrangebyscore(
                database_key, "+inf", "-inf", page, items_per_page, True
            ),
        )
        if database_key is not None
        else data.iloc[page : page + items_per_page - 1].items()
    )
    embed = discord.Embed(type="rich", colour=discord.Color.blurple())
    embed.set_author(name=config.options["bot_signature"])
    leaderboard = ""

    for i, stats in enumerate(leaderboard_list):
        leaderboard += f"{i+1+page}. **{stats[0]}** - {int(stats[1])}\n"
    embed.add_field(name=title, value=leaderboard, inline=False)

    await ctx.send(embed=embed)


def build_id_list(
    categories: Union[str, Iterable] = "", state: Union[str, Iterable] = ""
):
    """Generates an ID list based on given arguments

    - `categories`: category string/Iterable
    - `state`: state string/Iterable
    """
    logger.info("building id list")
    if isinstance(categories, str):
        categories = categories.split(" ")
    if isinstance(state, str):
        state = state.split(" ")

    id_choices = []
    group_args = set(
        map(dealias_group, all_categories.intersection(set(map(str.lower, categories))))
    )
    state_args = set(states.keys()).intersection(set(map(str.upper, state)))
    logger.info(f"group_args: {group_args}, state_args: {state_args}")

    if not config.options["id_groups"]:
        logger.info("no groups allowed")
        group_args = set()

    if group_args:
        items_in_group = set(
            itertools.chain.from_iterable(groups.get(o, []) for o in group_args)
        )
        if state_args:
            items_in_state = set(
                itertools.chain(*(states[state]["list"] for state in state_args))
            )
            id_choices = list(items_in_group.intersection(items_in_state))
        else:
            id_choices = list(items_in_group.intersection(set(id_list)))
    elif state_args:
        id_choices = list(
            set(itertools.chain(*(states[state]["list"] for state in state_args)))
        )
    else:
        id_choices = id_list

    logger.info(f"id_choices length: {len(id_choices)}")
    return id_choices


def backup_all():
    """Backs up the database to a file.

    This function serializes all data in the REDIS database
    into a file in the `backups` directory.

    This function is run with a task every 6 hours and sends the files
    to a specified discord channel.
    """
    logger.info("Starting Backup")
    logger.info("Creating Dump")
    keys = (key.decode("utf-8") for key in database.keys())
    dump = ((database.dump(key), key) for key in keys)
    logger.info("Finished Dump")
    logger.info("Writing To File")
    try:
        os.mkdir(config.options["backups_dir"])
        logger.info("Created backups directory")
    except FileExistsError:
        logger.info("Backups directory exists")
    with open(config.options["backups_dir"] + "dump.dump", "wb") as f:
        with open(config.options["backups_dir"] + "keys.txt", "w") as k:
            for item, key in dump:
                pickle.dump(item, f)
                k.write(f"{key}\n")
    logger.info("Backup Finished")


async def fools(ctx):
    logger.info(f"holiday check: invoked command: {str(ctx.command)}")
    if str(ctx.command) in ("leaderboard", "missed", "score", "streak", "userscore"):
        embed = discord.Embed(
            type="rich",
            colour=discord.Color.blurple(),
            title=f"{str(ctx.command).title()}",
        )
        embed.set_author(name=config.options["bot_signature"])
        embed.add_field(
            name=f"{str(ctx.command).title()}",
            value="User scores and data have been cleared. We apologize for the inconvenience.",
            inline=False,
        )
        await ctx.send(embed=embed)
        raise GenericError(code=666)
    return True


async def get_all_users(bot):
    logger.info("Starting user cache")
    user_ids = map(int, database.zrangebyscore("users:global", "-inf", "+inf"))
    for user_id in user_ids:
        await fetch_get_user(user_id, bot=bot, member=False)
    logger.info("User cache finished")


def evict_images():
    """Deletes images for items that have exceeded a certain frequency.

    This prevents images from being stale. If the item frequency has
    been incremented more than 10 times, this function will delete the
    top 3 items.
    """
    logger.info("Updating cached images")

    for item in map(
        lambda x: x.decode(),
        database.zrevrangebyscore(
            "frequency.item.refresh:global", "+inf", min=10, start=0, num=3
        ),
    ):
        database.zadd("frequency.item.refresh:global", {item: 0})
        category = get_category(item)
        path = f"{config.options['download_dir']}{category}/{item.lower()}/"
        if os.path.exists(path):
            shutil.rmtree(path)


class CustomCooldown:
    """Halve cooldown times in DM channels."""

    # Code adapted from discord.py example
    def __init__(
        self,
        per: float,
        disable: bool = False,
        bucket: commands.BucketType = commands.BucketType.channel,
    ):
        """Initialize a custom cooldown.

        `per` (float) - Cooldown default duration, halves in DM channels
        `bucket` (commands.BucketType) - cooldown scope, defaults to channel
        """
        rate = 1
        dm_per = per / 2
        race_per = 0.5
        self.disable = disable
        self.default_mapping = commands.CooldownMapping.from_cooldown(rate, per, bucket)
        self.dm_mapping = commands.CooldownMapping.from_cooldown(rate, dm_per, bucket)
        self.race_mapping = commands.CooldownMapping.from_cooldown(
            rate, race_per, bucket
        )

    def __call__(self, ctx: commands.Context):
        if not self.disable and ctx.guild is None:
            # halve cooldown in DMs
            bucket = self.dm_mapping.get_bucket(ctx.message)

        elif ctx.command.name.startswith("check") and ctx.channel.name.startswith(
            "racing"
        ):
            # tiny check cooldown in racing channels
            bucket = self.race_mapping.get_bucket(ctx.message)

        else:
            bucket = self.default_mapping.get_bucket(ctx.message)

        retry_after = bucket.update_rate_limit()
        if retry_after:
            raise commands.CommandOnCooldown(bucket, retry_after)
        return True
