#!/usr/bin/python3

from __future__ import print_function

import os
import re
import sys
import time
import logging
import logging
import optparse

import gratia.common.Gratia as Gratia
import gratia.services.StorageElement
import gratia.services.StorageElementRecord

log = logging.basicConfig(level=logging.INFO)

help = \
"""
**OSG Storage Space Reporting Tool**
This tool takes the status and description of a SE's storage space and reports
it to Gratia, the OSG's acccounting tool.  It is meant to use in conjunction
with other tools that interact and query a storage element.

A few explanations of the parameters:
  - uniqueID: This is the UniqueID of the SE; usually, it is the hostname of
    the primary SRM endpoint.  If there is no primary SRM endpoint, try to
    set this to the host name of the "primary" node in the system.  This must
    be a globally-unique string
  - Space Type: The type of the logical space.  Suggested strings are:
    - SpaceReservation: A SRM space reservation
    - Pool: A single disk server
    - PoolGroup: A group of disk servers
    - LinkGroup: A group of PoolGroups (dCache specific)
    - StaticSpaceReservation: A staticly configured space reservation (BeStMan)
    - Quota: A quota on the storage element.
    - Directory: All space used inside a specific directory.
    - SE: The entire Storage Element
    Note that some of these refer to raw disk used (Pool, PoolGroup) while
    others refer to logical usage within the storage element (Quota,
    SpaceReservation, Directory).
  - Space Name: An arbitrary name for the space.  Must be unique for the SE.
  - Space Unique ID.  This string is generated by the tool and has the
    following form:
      <SE UniqueID>:<space type>:<space name>
    This space name must be unique for the entire SE.
  - Parent ID.  This string (may be null) is set to the space unique ID for any
    logical parent that the space might have.  For example, a pool might have a
    parent pool group.  A pool group might have a parent link group.  A space
    reservation would have a parent link group. (All for dCache).  For many of
    the spaces, the parent space might be the entire storage element.  See the
    above comment about the Space Unique ID to know how to generate the
    Parent ID from a given space.
      - A Parent's space does not have to add up to the sum of its children!
        For example, a SE could have 10TB, and given 20 users with a 1TB quota.
  - VO: Any VO (or group) allowed to write into the space.  May be comma-
    seperated.  May be a wildcard.  May be left blank.  Not validated against
    a list of known VOs.
  - Owner: The owner of a space.  Might be a person's name (Johnny Public), a
    unix username, or a certificate's DN.
  - MeasurementType: Either "logical" or "raw"; refers to whether this is
    referring to space used on a hard drive, or the amount of space a user
    could use (taking replication into account, for example).
      
A few examples.  Say our SE is named "Example_SE", and is a dCache-based SE.

Suppose Johnny has a space reservation for 10GB, of which he has used 1GB.
Suppose also that Johnny is a member of CMS.  Then, a proper invocation of
this tool might be:

storageReport --config ProbeConfig --uniqueID srm.example.com \
  --se Example_SE --spaceName "CMSUSER_Johnny_12345" \
  --spaceType SpaceReservation --implementation dCache \
  --version 1.9.2-4 --status Production \
  --parent "srm.example.com:LinkGroup:cms_users" --vo CMS --owner Johnny \
  --measurementType logical --storageType disk \
  --total 10 --free 9 --used 1 --send
  
Suppose Example_SE is Hadoop-based, and we'd like to report space quotas usage
for Johnny; his quota is 10GB, and he has used 1GB

storageReport --config ProbeConfig --uniqueID srm.example.com \
  --se Example_SE --spaceName 'johnny_store_user' --spaceType Quota \
  --implementation Hadoop --version 0.19.1 \
  --parent "srm.example.com:SE:Example_SE" --vo CMS --owner johnny \
  --measurementType logical --storageType disk --total 10 --free 9 --used 1 \
  --send
  
Suppose that Example SE has a storage pool, dcache09.example.com, that is 40TB,
of which 30TB is used:

storageReport --config ProbeConfig --uniqueID srm.example.com \
  --se Example_SE --spaceName dcache09.example.com --spaceType Pool \
  --implementation Hadoop --version 0.19.1 \
  --parent "srm.example.com:SE:Example_SE" --measurementType raw \
  --storageType disk --total 40000 --free 10000 --used 30000 --send
  
"""

def build_opts():
    parser = optparse.OptionParser(usage=help)
    parser.add_option('-c', '--config', dest='ProbeConfig', type="string",
                      help='Location of the Gratia ProbeConfig file')
    parser.add_option('-i', '--uniqueID', dest='uniqueid', type="string",
                      help="The storage element's Unique ID (usually, the "
                      "hostname of the primary SRM endpoint")
    parser.add_option('-s', '--se', dest='SE', type="string",
                      help="The registered name of the SE")
    parser.add_option("-n", "--spaceName", dest="spaceName", type="string",
                      help="The name of the storage space")
    parser.add_option("-t", "--spaceType", dest="spaceType", type="string",
                      help="The type of the storage space")
    parser.add_option("-m", "--implementation", dest="implementation",
                      help="The implementation type of the SE", type="string")
    parser.add_option("-v", "--version", dest="version", type="string",
                      help="The version of the SE implementation")
    parser.add_option("-u", "--status", dest="status", type="string",
                      help="The status of the storage space [default: "
                      "%default].", default="Production")
    parser.add_option("-p", "--parent", dest="parentID", type="string",
                      help="The parent space unique ID.")
    parser.add_option("--vo", dest="vo", type="string",
                      help="VOs allowed to access this space [default: " \
                      "%default].", default="*")
    parser.add_option("-o", "--owner", dest="owner", type="string",
                      help="Owner of this space.")
    parser.add_option("--measurementType", dest="measurementType",
                      help="The type of the measurement (logical or raw) " \
                      "[default: %default].", type="string", default="raw")
    parser.add_option("--storageType", dest="storageType", type="string",
                      help="Type of underlying storage (disk or tape)" \
                      "[default: %default].", default="disk")
    parser.add_option("--total", dest="total", type="int",
                      help="Total space in GB")
    parser.add_option("--free", dest="free", type="int",
                      help="Free space in GB")
    parser.add_option("--used", dest="used", type="int",
                      help="Used space in GB")
    parser.add_option("--fileCountLimit", dest="fileCountLimit", type="int",
                      help="Maximum number of files.")
    parser.add_option("--fileCount", dest="fileCount", type="int",
                      help="Number of files in the space.")
    parser.add_option("--send", dest="send", action="store_true",
                      help="Set flag to send the report.")
    return parser

def printOpts(opt):
    str = ''
    str += '**Storage Report**\n'
    str += 'SE Unique ID: %s\n' % opt.uniqueid
    str += 'SE Name: %s\n' % opt.SE
    str += 'Space Name: %s\n' % opt.spaceName
    str += 'Space Type: %s\n' % opt.spaceType
    str += 'SE Implementation: %s\n' % opt.implementation
    str += 'SE Version: %s\n' % opt.version
    str += 'SE Status: %s\n' % opt.status
    str += 'Parent Space Unique ID: %s\n' % opt.parentID
    str += 'Space owner: %s\n' % opt.owner
    str += 'Measurement Type: %s\n' % opt.measurementType
    str += 'Type of underlying storage: %s\n' % opt.storageType
    str += 'Total space in GB: %s\n' % opt.total
    str += 'Free space in GB: %s\n' % opt.free
    str += 'Used space in GB: %s\n' % opt.used
    str += 'File count: %s\n' % opt.fileCount
    str += 'Maximum number of files: %s\n' % opt.fileCountLimit
    
    if opt.send:
        str += 'We will be sending this data to Gratia.\n'
    else:
        str += 'We will NOT be sending data to Gratia.\n'
        
    print(str)

valid_name = re.compile('[A-Za-z0-9_-]*')
valid_vo = re.compile('[A-Za-z0-9_\-\*]*')

requiredOpts = ['uniqueid', 'SE', 'spaceName', 'spaceType', 'implementation',
                'version', 'parentID', 'vo', 'owner', 'total', 'free', 'used']
requiredOptNames = ['SE Unique ID', 'SE', 'space name', 'space type',
                'implementation', 'version', 'parent space ID', 'VO', 
                'owner name', 'total (GB)', 'free (GB)', 'space (GB)']

def requireOpts(opt):
    for option, name in zip(requiredOpts, requiredOptNames):
        if getattr(opt, option) == None:
            raise Exception("Option for %s is not set!  This is a required" \
                            " option." % name)

def validateOpts(opt):
    requireOpts(opt)
    if not os.path.exists(opt.ProbeConfig):
        path = "/etc/gratia/storage-report/ProbeConfig"
        if os.path.exists(path):
            opt.ProbeConfig = path
    if not os.path.exists(opt.ProbeConfig):
        raise Exception("Could not find the ProbeConfig file in %s." % path)
    if not valid_name.match(opt.uniqueid):
        raise Exception("Invalid Unique ID: %s." % opt.uniqueid)
    if not valid_name.match(opt.SE):
        raise Exception("Invalid SE Name: %s." % opt.SE)
    if not valid_name.match(opt.spaceName):
        raise Exception("Invalid space name: %s." % opt.spaceName)
    if not valid_name.match(opt.spaceType):
        raise Exception("Invalid space type: %s." % opt.spaceType)
    if not valid_name.match(opt.implementation):
        raise Exception("Invalid implementation name: %s." % opt.implementation)
    names = str(opt.parentID).split(':')
    err_msg = "Invalid parent ID: %s.  A valid parent ID should be of the form" \
        " <SE unique ID>:<space type>:<space name>." % opt.parentID
    if len(names) != 3:
        raise Exception(err_msg)
    for name in names:
        if not valid_name.match(name):
            raise Exception(err_msg)
    if not valid_vo.match(opt.vo):
        raise Exception("Invalid VO name: %s." % opt.vo)
    opt.measurementType = str(opt.measurementType).lower()
    if opt.measurementType not in ['raw', 'logical']:
        raise Exception("Invalid measurement type: %s.  Valid options are "\
                        "'raw' or 'logical'." % opt.measurementType)
    opt.storageType = str(opt.storageType).lower()
    if opt.storageType not in ['disk', 'tape']:
        raise Exception("Invalid storage type: %s.  Valid options are 'disk'"\
                        " or 'tape'")
    try:
        opt.total = int(opt.total)
    except:
        raise Exception("Unable to convert total %s to integer." % opt.total)
    try:
        opt.free = int(opt.free)
    except:
        raise Exception("Unable to convert free %s to integer." % opt.free)
    try:
        opt.used = int(opt.used)
    except:
        raise Exception("Unable to convert used %s to integer." % opt.used)
    
def sendReport(opt):
    Gratia.Initialize(opt.ProbeConfig)
    t = time.time()
    desc = StorageElement.StorageElement()
    state = StorageElementRecord.StorageElementRecord()
    desc.Timestamp(t)
    state.Timestamp(t)
    uniqueID = '%s:%s:%s' % (opt.uniqueid, opt.spaceType, opt.spaceName)
    desc.UniqueID(uniqueID)
    desc.Name(opt.spaceName)
    desc.SE(opt.SE)
    desc.SpaceType(opt.spaceType)
    desc.Implementation(opt.implementation)
    desc.Version(opt.version)
    desc.Status(opt.status)
    desc.ParentID(opt.parentID)
    desc.VO(opt.vo)
    desc.OwnerDN(opt.owner)
    state.UniqueID(uniqueID)
    state.MeasurementType(opt.measurementType)
    state.StorageType(opt.storageType)
    state.TotalSpace(opt.total*1000**3)
    state.FreeSpace(opt.free*1000**3)
    state.UsedSpace(opt.used*1000**3)
    if opt.fileCount != None:
        state.FileCount(opt.fileCount)
    if opt.fileCountLimit != None:
        state.FileCountLimit(opt.fileCountLimit)

    print("SE Description to Gratia", Gratia.Send(desc))
    print("SE State to Gratia", Gratia.Send(state))
    
def main():
    parser = build_opts()
    options, args = parser.parse_args()
    validateOpts(options)
    printOpts(options)
    if options.send:
        sendReport(options)
    
if __name__ == '__main__':
    main()
