#!/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 rse as client
    from rucio.api import config
    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 TB(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
    try:
        spacetoken = client.list_rse_attributes(rse)['spacetoken']
    except:
        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, TB(capacity), TB(storageused), TB(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" % (TB(used), endpoint['name'])
            used_others += used

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

    minfree = min(capacity * policy[token][0] / 100.0, policy[token][1] * (1000**4))
    print "RSE %s: calculated minimum free space %sTB" % (rse, TB(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, TB(minfree))

sys.exit(OK)
