# Copyright 2013-2018 CERN for the benefit of the ATLAS collaboration.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors:
#  - Martin Barisits, <martin.barisits@cern.ch>, 2013-2018
#  - Cedric Serfon, <cedric.serfon@cern.ch>, 2015
#  - Vincent Garonne, <vincent.garonne@cern.ch>, 2015
#  - Hannes Hansen, <hannes.jakob.hansen@cern.ch>, 2018-2019
#  - Andrew Lister, <andrew.lister@stfc.ac.uk>, 2019
#  - Patrick Austin, <patrick.austin@stfc.ac.uk>, 2020
#  - Eli Chadwick, <eli.chadwick@stfc.ac.uk>, 2020
#
# PY3K COMPATIBLE

from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.sql import func
from sqlalchemy.sql.expression import and_, or_

from rucio.core.account import get_all_rse_usages_per_account
from rucio.core.rse import get_rse_name
from rucio.core.rse_expression_parser import parse_expression
from rucio.db.sqla import models
from rucio.db.sqla.session import read_session, transactional_session


@read_session
def get_rse_account_usage(rse_id, session=None):
    """
    Returns the account limit and usage for all for all accounts on a RSE.

    :param rse_id:   The id of the RSE.
    :param session:  Database session in use.
    :return:         List of dictionaries.
    """
    result = []
    query = session.query(models.AccountUsage.account, models.AccountUsage.files, models.AccountUsage.bytes, models.AccountLimit.bytes, models.RSE.rse)
    query = query.filter(models.RSE.id == models.AccountUsage.rse_id)
    query = query.outerjoin(models.AccountLimit, and_(models.AccountUsage.account == models.AccountLimit.account, models.AccountUsage.rse_id == models.AccountLimit.rse_id)).filter(models.AccountUsage.rse_id == rse_id)
    account_limits_tmp = query.all()
    for row in account_limits_tmp:
        result.append({'rse_id': rse_id, 'rse': row[4], 'account': row[0], 'used_files': row[1], 'used_bytes': row[2], 'quota_bytes': row[3]})
    return result


@read_session
def get_global_account_limits(account=None, session=None):
    """
    Returns the global account limits for the account.

    :param account:  Account to check the limit for.
    :param session:  Database session in use.
    :return:         Dict {'MOCK': {'resolved_rses': ['MOCK'], 'limit': 10, 'resolved_rse_ids': [123]}}.
    """
    if account:
        global_account_limits = session.query(models.AccountGlobalLimit).filter_by(account=account).all()
    else:
        global_account_limits = session.query(models.AccountGlobalLimit).all()

    resolved_global_account_limits = {}
    for limit in global_account_limits:
        if account:
            resolved_rses = parse_expression(limit['rse_expression'], filter={'vo': account.vo}, session=session)
        else:
            resolved_rses = parse_expression(limit['rse_expression'], session=session)
        limit_in_bytes = limit['bytes']
        if limit_in_bytes == -1:
            limit_in_bytes = float('inf')
        resolved_global_account_limits[limit['rse_expression']] = {
            'resolved_rses': [resolved_rse['rse'] for resolved_rse in resolved_rses],
            'resolved_rse_ids': [resolved_rse['id'] for resolved_rse in resolved_rses],
            'limit': limit_in_bytes
        }
    return resolved_global_account_limits


@read_session
def get_global_account_limit(account, rse_expression, session=None):
    """
    Returns the global account limit for the account on the rse expression.

    :param account:         Account to check the limit for.
    :param rse_expression:  RSE expression to check the limit for.
    :param session:         Database session in use.
    :return:                Limit in Bytes.
    """
    try:
        global_account_limit = session.query(models.AccountGlobalLimit).filter_by(account=account, rse_expression=rse_expression).one()
        if global_account_limit.bytes == -1:
            return float("inf")
        else:
            return global_account_limit.bytes
    except NoResultFound:
        return None


@read_session
def get_local_account_limit(account, rse_id, session=None):
    """
    Returns the account limit for the account on the rse.

    :param account:  Account to check the limit for.
    :param rse_id:   RSE id to check the limit for.
    :param session:  Database session in use.
    :return:         Limit in Bytes.
    """
    try:
        account_limit = session.query(models.AccountLimit).filter(models.AccountLimit.account == account,
                                                                  models.AccountLimit.rse_id == rse_id).one()
        if account_limit.bytes == -1:
            return float("inf")
        else:
            return account_limit.bytes
    except NoResultFound:
        return None


@read_session
def get_local_account_limits(account, rse_ids=None, session=None):
    """
    Returns the account limits for the account on the list of rses.

    :param account:  Account to check the limit for.
    :param rse_ids:  List of RSE ids to check the limit for.
    :param session:  Database session in use.
    :return:         Dictionary {'rse_id': bytes, ...}.
    """

    account_limits = {}
    if rse_ids:
        rse_id_clauses = []
        for rse_id in rse_ids:
            rse_id_clauses.append(and_(models.AccountLimit.rse_id == rse_id, models.AccountLimit.account == account))
        rse_id_clause_chunks = [rse_id_clauses[x:x + 10] for x in range(0, len(rse_id_clauses), 10)]
        for rse_id_chunk in rse_id_clause_chunks:
            tmp_limits = session.query(models.AccountLimit).filter(or_(*rse_id_chunk)).all()
            for limit in tmp_limits:
                if limit.bytes == -1:
                    account_limits[limit.rse_id] = float("inf")
                else:
                    account_limits[limit.rse_id] = limit.bytes
    else:
        account_limits_tmp = session.query(models.AccountLimit).filter(models.AccountLimit.account == account).all()
        for limit in account_limits_tmp:
            if limit.bytes == -1:
                account_limits[limit.rse_id] = float("inf")
            else:
                account_limits[limit.rse_id] = limit.bytes
    return account_limits


@transactional_session
def set_local_account_limit(account, rse_id, bytes, session=None):
    """
    Returns the limits for the account on the rse.

    :param account:  Account to check the limit for.
    :param rse_id:   RSE id to check the limit for.
    :param bytes:    The limit value, in bytes, to set.
    :param session:  Database session in use.
    """
    try:
        account_limit = session.query(models.AccountLimit).filter(models.AccountLimit.account == account,
                                                                  models.AccountLimit.rse_id == rse_id).one()
        account_limit.bytes = bytes
    except NoResultFound:
        models.AccountLimit(account=account, rse_id=rse_id, bytes=bytes).save(session=session)


@transactional_session
def set_global_account_limit(account, rse_expression, bytes, session=None):
    """
    Sets the global limit for the account on a RSE expression.

    :param account:         Account to check the limit for.
    :param rse_expression:  RSE expression to check the limit for.
    :param bytes:           The limit value, in bytes, to set.
    :param session:         Database session in use.
    """
    try:
        account_limit = session.query(models.AccountGlobalLimit).filter(models.AccountGlobalLimit.account == account,
                                                                        models.AccountGlobalLimit.rse_expression == rse_expression).one()
        account_limit.bytes = bytes
    except NoResultFound:
        models.AccountGlobalLimit(account=account, rse_expression=rse_expression, bytes=bytes).save(session=session)


@transactional_session
def delete_local_account_limit(account, rse_id, session=None):
    """
    Deletes a local account limit.

    :param account:  Account to delete the limit for.
    :param rse_id:   RSE id to delete the limit for.
    :param session:  Database session in use.
    :returns:        True if something was deleted; False otherwise.
    """
    try:
        session.query(models.AccountLimit).filter(models.AccountLimit.account == account,
                                                  models.AccountLimit.rse_id == rse_id).one().delete(session=session)
        return True
    except NoResultFound:
        return False


@transactional_session
def delete_global_account_limit(account, rse_expression, session=None):
    """
    Deletes a global account limit.

    :param account:         Account to delete the limit for.
    :param rse_expression:  RSE expression to delete the limit for.
    :param session:         Database session in use.
    :returns:               True if something was deleted; False otherwise.
    """
    try:
        session.query(models.AccountGlobalLimit).filter(models.AccountGlobalLimit.account == account,
                                                        models.AccountGlobalLimit.rse_expression == rse_expression).one().delete(session=session)
        return True
    except NoResultFound:
        return False


@transactional_session
def get_local_account_usage(account, rse_id=None, session=None):
    """
    Read the account usage and connect it with (if available) the account limits of the account.

    :param account:  The account to read.
    :param rse_id:   The rse_id to read (If none, get all).
    :param session:  Database session in use.

    :returns:        List of dicts {'rse_id', 'bytes_used', 'files_used', 'bytes_limit'}
    """

    if not rse_id:
        # All RSESs
        limits = get_local_account_limits(account=account, session=session)
        counters = session.query(models.AccountUsage).filter_by(account=account).all()
    else:
        # One RSE
        limits = get_local_account_limits(account=account, rse_ids=[rse_id], session=session)
        counters = session.query(models.AccountUsage).filter_by(account=account, rse_id=rse_id).all()
    result_list = []

    for counter in counters:
        if counter.bytes > 0 or counter.files > 0 or rse_id in limits.keys():
            result_list.append({'rse_id': counter.rse_id, 'rse': get_rse_name(rse_id=counter.rse_id, session=session),
                                'bytes': counter.bytes, 'files': counter.files,
                                'bytes_limit': limits.get(counter.rse_id, 0),
                                'bytes_remaining': limits.get(counter.rse_id, 0) - counter.bytes})
    return result_list


@transactional_session
def get_global_account_usage(account, rse_expression=None, session=None):
    """
    Read the account usage and connect it with the global account limits of the account.

    :param account:          The account to read.
    :param rse_expression:   The RSE expression (If none, get all).
    :param session:          Database session in use.

    :returns:                List of dicts {'rse_id', 'bytes_used', 'files_used', 'bytes_limit'}
    """
    result_list = []
    if not rse_expression:
        # All RSE Expressions
        limits = get_global_account_limits(account=account, session=session)
        all_rse_usages = {usage['rse_id']: (usage['bytes'], usage['files']) for usage in get_all_rse_usages_per_account(account=account, session=session)}
        for rse_expression, limit in limits.items():
            usage = 0
            files = 0
            for rse in limit['resolved_rse_ids']:
                usage += all_rse_usages.get(rse, [0])[0]
                files += all_rse_usages.get(rse, [0, 0])[1]
            result_list.append({'rse_expression': rse_expression,
                                'bytes': usage, 'files': files,
                                'bytes_limit': limit['limit'],
                                'bytes_remaining': limit['limit'] - usage})
    else:
        # One RSE Expression
        limit = get_global_account_limit(account=account, rse_expression=rse_expression, session=session)
        vo = account.vo
        resolved_rses = [resolved_rse['id'] for resolved_rse in parse_expression(rse_expression, filter={'vo': vo}, session=session)]
        usage = session.query(func.sum(models.AccountUsage.bytes), func.sum(models.AccountUsage.files))\
                       .filter(models.AccountUsage.account == account, models.AccountUsage.rse_id.in_(resolved_rses))\
                       .group_by(models.AccountUsage.account).first()
        result_list.append({'rse_expression': rse_expression,
                            'bytes': usage[0], 'files': usage[1],
                            'bytes_limit': limit,
                            'bytes_remaining': limit - usage[0]})
    return result_list
