#!/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
#

# 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 SRM 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 - SRM free, where SRM
# free is the total SRM 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
import requests
from urlparse import urlparse

# Try to use server environment (direct database access). If that fails use
# client.
server = False
try:
    from rucio.api import rse as c
    from rucio.api import config
    server = True
except:
    from rucio.client import Client
    from rucio.client.configclient import ConfigClient
    c = 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'
    resp = requests.get(url=url)
    data = json.loads(resp.content)
except Exception, e:
    print "Failed to get information from AGIS: %s" % str(e)
    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 c.list_rses()]
except Exception, e:
    print "Failed to get RSEs from Rucio: %s" % str(e)
    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, e:
    print "Failed to get configuration information from Rucio: %s" % str(e)
    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 = c.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:
            capacity = c.get_rse_usage(rse, None, source='srm')
        else:
            capacity = c.get_rse_usage(rse, filters={'source': 'srm'})
    except Exception, e:
        print "Could not get SRM information for %s: %s" % (rse, str(e))
        continue

    capacity = [source['total'] for source in capacity if source['source'] == 'srm']
    if not capacity:
        print 'No SRM information available for %s' % rse
        continue
    capacity = capacity[0]

    print "RSE %s: total capacity %sTB" % (rse, TB(capacity))

    # 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 = c.get_rse_usage(endpoint['name'], None, source='rucio')
                else:
                    used = c.get_rse_usage(endpoint['name'], filters={'source': 'rucio'})
            except Exception, e:
                print "Could not get used Rucio space for %s: %s" % (endpoint['name'], str(e))
                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))

    # Now add the space used in other tokens to the limit
    minfree += used_others
    try:
        if server:
            c.set_rse_limits(rse, 'MinFreeSpace', minfree, 'root')
            c.set_rse_limits(rse, 'MaxBeingDeletedFiles', max_files_to_delete, 'root')
        else:
            c.set_rse_limits(rse, 'MinFreeSpace', minfree)
            c.set_rse_limits(rse, 'MaxBeingDeletedFiles', max_files_to_delete)
    except Exception, e:
        print "Failed to set RSE limits for %s: %s" % (rse, str(e))
        continue

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

sys.exit(OK)
