#!/usr/bin/env python
"""
 Copyright European Organization for Nuclear Research (CERN) 2014

 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

 Authors:
 - David Cameron, <david.cameron@cern.ch>, 2014-2016
 - Mario Lassnig, <mario.lassnig@cern.ch>, 2015
 - Cedric Serfon, <cedric.serfon@cern.ch>, 2017

 Sets the minimum free space on RSEs according to the policy, which is set in
 the configuration table of Rucio server. A relative and absolute limit are
 set for the relevant endpoints, for example:
  Spacetoken  Free ratio  Free absolute
  PRODDISK      25%         10.0 TB

 The smaller of ratio and absolute is the threshold below which to clean.
 Some tokens (tape, groupdisk, localgroupdisk) are not cleaned automatically.

 The capacity of each RSE is Storage used - Rucio used of other RSEs sharing the
 token. This allows RSEs to use space pledged but not used by other RSEs. The
 minimum free space is evaluated based on this capacity. In the reaper Rucio
 calculates the space to clean as MinFreeSpace limit - Storage free, where Storage
 free is the total Storage capacity - Rucio used for this RSE. Therefore the
 MinFreeSpace limit set here must include all the used space for all the other
 RSEs in the token.
"""

import json
import sys
from urlparse import urlparse
import requests

# Try to use server environment (direct database access). If that fails use
# client.
server = False
try:
    from rucio.api import config
    from rucio.api import rse as client
    server = True
except:
    from rucio.client import Client
    from rucio.client.configclient import ConfigClient
    client = Client()
    config = ConfigClient()

UNKNOWN, OK, WARNING, CRITICAL = -1, 0, 1, 2

# This is the limit of files to delete in each RSE in the reaper loop. To be
# decided what is the ideal value and if it should be per RSE.
max_files_to_delete = 100


def toTB(size):
    return size / 1000.0**4


# Get endpoint info from AGIS to know the RSEs in each space token
try:
    URL = 'http://atlas-agis-api.cern.ch/request/ddmendpoint/query/list/?json&state=ACTIVE&site_state=ACTIVE'
    RESP = requests.get(url=URL)
    DATA = json.loads(RESP.content)
except Exception as error:
    print "Failed to get information from AGIS: %s" % str(error)
    sys.exit(CRITICAL)

# Map of RSE: hostname
RSE_HOST = {}
for endpoint in DATA:
    host = urlparse(endpoint['se']).hostname
    if host:
        RSE_HOST[endpoint['name']] = host

try:
    RSES = [rse['rse'] for rse in client.list_rses()]
except Exception as error:
    print "Failed to get RSEs from Rucio: %s" % str(error)
    sys.exit(CRITICAL)

# Get policy defined in config. Each section is called limitstoken
# {token: (relative limit in %, absolute limit in TB)}
policy = {}
try:
    if server:
        SECTIONS = config.sections('root')
        for section in [s for s in SECTIONS if s.startswith('limits')]:
            policy[section[6:].upper()] = (config.get(section, 'rellimit', 'root'), config.get(section, 'abslimit', 'root'))
    else:
        SECTIONS = config.get_config()
        for section in [s for s in SECTIONS if s.startswith('limits')]:
            policy[section[6:].upper()] = (SECTIONS[section]['rellimit'], SECTIONS[section]['abslimit'])

except Exception as error:
    print "Failed to get configuration information from Rucio: %s" % str(error)
    sys.exit(CRITICAL)

for rse in RSES:

    tokens = [token for token in policy if rse.endswith(token)]
    if not tokens:
        continue

    if len(tokens) != 1:
        print "RSE %s has multiple limits defined" % rse
        continue

    token = tokens[0]

    if not [r for r in DATA if r['name'] == rse]:
        print "RSE %s not defined in AGIS" % rse
        continue
    rse_attr = client.list_rse_attributes(rse)
    if 'spacetoken' in rse_attr:
        spacetoken = client.list_rse_attributes(rse)['spacetoken']
    else:
        print "No space token info for %s" % rse
        continue

    # Client and server API are different for get_rse_usage
    try:
        if server:
            spaceinfo = client.get_rse_usage(rse, None)
        else:
            spaceinfo = client.get_rse_usage(rse)
    except Exception as error:
        print "Could not get space information for %s: %s" % (rse, str(error))
        continue

    spaceinfo = [i for i in spaceinfo]  # Generator can only be used once

    capacity = [source['total'] for source in spaceinfo if source['source'] == 'storage']
    storageused = [source['used'] for source in spaceinfo if source['source'] == 'storage']
    rucioused = [source['used'] for source in spaceinfo if source['source'] == 'rucio']
    if not capacity or not storageused or not rucioused:
        print 'Missing space info for %s' % rse
        continue
    capacity = capacity[0]
    storageused = storageused[0]
    rucioused = rucioused[0]

    print "RSE %s: total capacity %sTB, Storage used %sTB, Rucio used %sTB" % (rse, toTB(capacity), toTB(storageused), toTB(rucioused))

    # If this RSE shares space with others remove rucio used from total space
    # to calculate the limit
    used_others = 0
    for endpoint in DATA:
        if endpoint['name'] != rse and (RSE_HOST[endpoint['name']] == RSE_HOST[rse] and spacetoken == endpoint['token']):
            try:
                if server:
                    used = client.get_rse_usage(endpoint['name'], None, source='rucio')
                else:
                    used = client.get_rse_usage(endpoint['name'], filters={'source': 'rucio'})
            except Exception as error:
                print "Could not get used Rucio space for %s: %s" % (endpoint['name'], str(error))
                continue

            used = [source['used'] for source in used if source['source'] == 'rucio']
            if not used:
                print "No Rucio used space information for %s" % rse
                continue
            used = used[0]

            print "Removing %fTB used space in %s" % (toTB(used), endpoint['name'])
            used_others += used

    capacity -= used_others
    print "Remaining capacity for %s: %sTB" % (rse, toTB(capacity))

    if 'MinFreeSpace' in rse_attr:
        minfree = int(rse_attr['MinFreeSpace'])
        print "RSE %s: Will apply forced value for minimum free space %sTB" % (rse, toTB(minfree))
    else:
        minfree = min(capacity * policy[token][0] / 100.0, policy[token][1] * (1000**4))
        print "RSE %s: calculated minimum free space %sTB" % (rse, toTB(minfree))

    try:
        if server:
            client.set_rse_limits(rse, 'MinFreeSpace', minfree, 'root')
            client.set_rse_limits(rse, 'MaxBeingDeletedFiles', max_files_to_delete, 'root')
        else:
            client.set_rse_limits(rse, 'MinFreeSpace', minfree)
            client.set_rse_limits(rse, 'MaxBeingDeletedFiles', max_files_to_delete)
    except Exception as error:
        print "Failed to set RSE limits for %s: %s" % (rse, str(error))
        continue

    print "Set MinFreeSpace for %s to %fTB" % (rse, toTB(minfree))

sys.exit(OK)
