#!/usr/bin/python3
import datetime

from gratia.common.Gratia import DebugPrint
#import gratia.common.GratiaWrapper as GratiaWrapper
import gratia.common.Gratia as Gratia

from gratia.common2.meter import GratiaMeter
from gratia.common2.pgpinput import PgInput
from gratia.common2.checkpoint import DateTransactionAuxCheckpoint
import gratia.common2.timeutil as timeutil


class _EnstoreTapeDriveInputStub:
    """Stub class, needs to be defined before the regular one, to avoid NameError

    Query producing the data in the test array:
::
Query: accounting=> select * from tape_mounts where start > '2014-11-15';
       node        | volume |    type     |   logname   |        start        |       finish        | state | storage_group | reads | writes
-------------------+--------+-------------+-------------+---------------------+---------------------+-------+---------------+-------+--------
 enmvr021.fnal.gov | TST082 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 12:40:24 | 2014-11-21 12:41:13 | M     | ANM           |     0 |      0
 enmvr022.fnal.gov | TST066 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 12:42:31 | 2014-11-21 12:43:10 | M     | ANM           |     0 |      0
 enmvr021.fnal.gov | TST082 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 12:44:23 | 2014-11-21 12:44:44 | D     | ANM           |     0 |      6
 enmvr021.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 12:44:45 | 2014-11-21 12:45:20 | M     | ANM           |     0 |      0
 enmvr022.fnal.gov | TST066 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 12:45:21 | 2014-11-21 12:45:44 | D     | ANM           |     1 |      9
 enmvr022.fnal.gov | TST070 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 12:52:35 | 2014-11-21 12:53:07 | M     | ANM           |     0 |      0
 enmvr021.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 12:57:29 | 2014-11-21 12:57:51 | D     | ANM           |    17 |     12
 enmvr021.fnal.gov | TST066 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 12:57:52 | 2014-11-21 12:58:44 | M     | ANM           |     0 |      0
 enmvr022.fnal.gov | TST070 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 12:59:21 | 2014-11-21 12:59:45 | D     | ANM           |     3 |      3
 enmvr022.fnal.gov | TST082 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 12:59:45 | 2014-11-21 13:00:30 | M     | ANM           |     0 |      0
 enmvr021.fnal.gov | TST066 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 13:00:53 | 2014-11-21 13:01:17 | D     | ANM           |    20 |      0
 enmvr021.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 13:01:28 | 2014-11-21 13:02:14 | M     | ANM           |     0 |      0
 enmvr022.fnal.gov | TST082 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 13:03:15 | 2014-11-21 13:03:29 | D     | ANM           |     1 |      0
 enmvr022.fnal.gov | TST063 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 13:03:31 | 2014-11-21 13:04:22 | M     | ANM           |     0 |      0
 enmvr021.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 13:04:54 | 2014-11-21 13:05:13 | D     | ANM           |     1 |      2
 enmvr021.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 13:05:37 | 2014-11-21 13:06:09 | M     | ANM           |     0 |      0
 enmvr022.fnal.gov | TST063 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 13:07:41 | 2014-11-21 13:08:01 | D     | ANM           |     0 |      1
 enmvr021.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 13:07:55 | 2014-11-21 13:08:14 | D     | ANM           |     1 |      0
 enmvr022.fnal.gov | TST063 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 13:08:02 | 2014-11-21 13:08:32 | M     | ANM           |     0 |      0
 enmvr021.fnal.gov | TST066 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 13:09:53 | 2014-11-21 13:10:31 | M     | ANM           |     0 |      0
 enmvr022.fnal.gov | TST063 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 13:11:50 | 2014-11-21 13:12:10 | D     | ANM           |     2 |      0
 enmvr021.fnal.gov | TST066 | ULTRIUM-TD4 | LTO4_021BMV | 2014-11-21 13:12:06 | 2014-11-21 13:12:27 | D     | ANM           |     1 |      0
 enmvr022.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 13:12:11 | 2014-11-21 13:13:16 | M     | ANM           |     0 |      0
 enmvr022.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_022MV  | 2014-11-21 13:15:13 | 2014-11-21 13:15:33 | D     | ANM           |     0 |      2
 enmvr021.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_021BMV | 2014-12-01 11:51:39 | 2014-12-01 11:52:33 | M     | ANM           |     0 |      0
 enmvr021.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_021BMV | 2014-12-01 11:54:27 | 2014-12-01 11:54:46 | D     | ANM           |     0 |      1
 enmvr022.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_022MV  | 2014-12-01 12:04:49 | 2014-12-01 12:05:45 | M     | ANM           |     0 |      0
 enmvr022.fnal.gov | TST065 | ULTRIUM-TD4 | LTO4_022MV  | 2014-12-01 12:07:39 | 2014-12-01 12:07:53 | D     | ANM           |     0 |      1

"""

    # All: node,volume,type,logname,start,finish,state,storage_group,reads,writes
    value_matrix = [['enmvr021.fnal.gov', 'TST082', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 12:40:24', '2014-11-21 12:41:13', 'M', 'ANM', 0, 0],
                    ['enmvr022.fnal.gov', 'TST066', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 12:42:31', '2014-11-21 12:43:10', 'M', 'ANM', 0, 0],
                    ['enmvr021.fnal.gov', 'TST082', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 12:44:23', '2014-11-21 12:44:44', 'D', 'ANM', 0, 6],
                    ['enmvr021.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 12:44:45', '2014-11-21 12:45:20', 'M', 'ANM', 0, 0],
                    ['enmvr022.fnal.gov', 'TST066', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 12:45:21', '2014-11-21 12:45:44', 'D', 'ANM', 1, 9],
                    ['enmvr022.fnal.gov', 'TST070', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 12:52:35', '2014-11-21 12:53:07', 'M', 'ANM', 0, 0],
                    ['enmvr021.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 12:57:29', '2014-11-21 12:57:51', 'D', 'ANM', 17, 12],
                    ['enmvr021.fnal.gov', 'TST066', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 12:57:52', '2014-11-21 12:58:44', 'M', 'ANM', 0, 0],
                    ['enmvr022.fnal.gov', 'TST070', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 12:59:21', '2014-11-21 12:59:45', 'D', 'ANM', 3, 3],
                    ['enmvr022.fnal.gov', 'TST082', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 12:59:45', '2014-11-21 13:00:30', 'M', 'ANM', 0, 0],
                    ['enmvr021.fnal.gov', 'TST066', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 13:00:53', '2014-11-21 13:01:17', 'D', 'ANM', 20, 0],
                    ['enmvr021.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 13:01:28', '2014-11-21 13:02:14', 'M', 'ANM', 0, 0],
                    ['enmvr022.fnal.gov', 'TST082', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 13:03:15', '2014-11-21 13:03:29', 'D', 'ANM', 1, 0],
                    ['enmvr022.fnal.gov', 'TST063', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 13:03:31', '2014-11-21 13:04:22', 'M', 'ANM', 0, 0],
                    ['enmvr021.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 13:04:54', '2014-11-21 13:05:13', 'D', 'ANM', 1, 2],
                    ['enmvr021.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 13:05:37', '2014-11-21 13:06:09', 'M', 'ANM', 0, 0],
                    ['enmvr022.fnal.gov', 'TST063', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 13:07:41', '2014-11-21 13:08:01', 'D', 'ANM', 0, 1],
                    ['enmvr021.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 13:07:55', '2014-11-21 13:08:14', 'D', 'ANM', 1, 0],
                    ['enmvr022.fnal.gov', 'TST063', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 13:08:02', '2014-11-21 13:08:32', 'M', 'ANM', 0, 0],
                    ['enmvr021.fnal.gov', 'TST066', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 13:09:53', '2014-11-21 13:10:31', 'M', 'ANM', 0, 0],
                    ['enmvr022.fnal.gov', 'TST063', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 13:11:50', '2014-11-21 13:12:10', 'D', 'ANM', 2, 0],
                    ['enmvr021.fnal.gov', 'TST066', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-11-21 13:12:06', '2014-11-21 13:12:27', 'D', 'ANM', 1, 0],
                    ['enmvr022.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 13:12:11', '2014-11-21 13:13:16', 'M', 'ANM', 0, 0],
                    ['enmvr022.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-11-21 13:15:13', '2014-11-21 13:15:33', 'D', 'ANM', 0, 2],
                    ['enmvr021.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-12-01 11:51:39', '2014-12-01 11:52:33', 'M', 'ANM', 0, 0],
                    ['enmvr021.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_021BMV', '2014-12-01 11:54:27', '2014-12-01 11:54:46', 'D', 'ANM', 0, 1],
                    ['enmvr022.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-12-01 12:04:49', '2014-12-01 12:05:45', 'M', 'ANM', 0, 0],
                    ['enmvr022.fnal.gov', 'TST065', 'ULTRIUM-TD4', 'LTO4_022MV', '2014-12-01 12:07:39', '2014-12-01 12:07:53', 'D', 'ANM', 0, 1]
                    ]

    def get_records():
        for i in _EnstoreTapeDriveInputStub.value_matrix:
            # All: node,volume,type,logname,start,finish,state,storage_group,reads,writes
            retv = {'node': i[0],
                    'volume': i[1],
                    'type': i[2],
                    'logname': i[3],
                    'start': timeutil.parse_datetime(i[4]),
                    'finish': timeutil.parse_datetime(i[5]),
                    'state': i[6],
                    'storage_group': i[7],
                    'reads': int(i[8]),
                    'writes': int(i[9]),
                    }
            yield retv
    get_records = staticmethod(get_records)


class EnstoreTapeDriveInput(PgInput):
    """Get tape drive usage information from the Enstore accounting DB.

    The table has no monotone unique ID.
    Uses DateTransactionCheckpoint and stores in the transaction field the number of seconds between the
    last unhandled mount (checkpoint date) and the last dismount. Date include the extremes not to loose record when
    they are added with the timestamp equal to the last one.
    """

    VERSION_ATTRIBUTE = 'EnstoreVersion'
    OK_TO_SEND_RECORD = {'state': 'oktosend', 'volume': 'oktosend'}

    def __init__(self, conn=None):
        PgInput.__init__(self, conn)
        self.input_min_interval = 0
        self.input_delay = 0
        self.rollback = 0

    def get_init_params(self):
        """Return list of parameters to read form the config file"""
        return PgInput.get_init_params(self) + [EnstoreTapeDriveInput.VERSION_ATTRIBUTE,
                'InputMinInterval', 'InputDelay', 'CheckpointRollback']

    def add_checkpoint(self, fname=None, max_val=None, default_val=None, fullname=False):
        """Add a checkpoint, default file name is cfp-INPUT_NAME

        :param fname: checkpoint file name (considered as prefix unless fullname=True)
                file name is fname-INPUT_NAME
        :param max_val: trim value for the checkpoint
        :param default_val: value if no checkpoint is available
        :param fullname: Default: False, if true, fname is considered the full file name
        :return:
        """

        if not fname:
            fname = "cpf-%s" % self.get_name()
        else:
            if not fullname:
                fname = "%s-%s" % (fname, self.get_name())
        if max_val is not None or default_val is not None:
            self.checkpoint = DateTransactionAuxCheckpoint(fname, max_val, default_val)
        else:
            self.checkpoint = DateTransactionAuxCheckpoint(fname)

    def start(self, static_info):
        """open DB connection and set version form config file"""
        PgInput.start(self, static_info)
        DebugPrint(4, "ETDI start, static info: %s" % static_info)
        if EnstoreTapeDriveInput.VERSION_ATTRIBUTE in static_info:
            self._set_version_config(static_info[EnstoreTapeDriveInput.VERSION_ATTRIBUTE])
        try:
            self.input_min_interval = int(static_info['InputMinInterval'])
        except (KeyError, ValueError):
            pass
        try:
            self.input_delay = int(static_info['InputDelay'])
        except (KeyError, ValueError):
            pass
        try:
            self.rollback = int(static_info['CheckpointRollback'])
        except (KeyError, ValueError):
            pass

    def _start_stub(self, static_info):
        """start replacement for testing: database connection errors are trapped"""
        try:
            DebugPrint(4, "Testing DB connection. The probe will not use it")
            PgInput.start(self, static_info)
            if self.status_ok():
                DebugPrint(4, "Connection successful")
            else:
                DebugPrint(4, "Connection failed")
            DebugPrint(4, "Closing the connection")
            self.stop()
        except:
            DebugPrint(1, "Database connection failed. The test can continue since stubs are used.")
        DebugPrint(4, "ETDI start stub, static info: %s" % static_info)
        if EnstoreTapeDriveInput.VERSION_ATTRIBUTE in static_info:
            self._set_version_config(static_info[EnstoreTapeDriveInput.VERSION_ATTRIBUTE])

    def get_version(self):
        # RPM package is 'enstore'
        return self._get_version('enstore')

    @staticmethod
    def get_dismount_mount_record(r, start_time):
        """Provide a mount record to trigger the dismount of a previous mount record r.

        It is used if no record on the same tape is found after the limit is reached.
        This record will never trigger a Gratia record since it is the last one sent to the meter
        and no other record will follow on the same volume.

        :param r: mount record to dismount
        :param start_time: tart time of the new mount record
        :return: mount record (dictionary with DB record keys)
        """
        retv = {'state': 'M',
                'node': r['node'],
                'volume': r['volume'],
                'type': r['type'],
                'logname': r['logname'],
                'start': start_time,
                'finish': timeutil.wind_time(start_time, seconds=20, backward=False),
                'storage_group': 'PROBE_UTILITY',
                'reads': 0,
                'writes':0
                }
        return retv

    @staticmethod
    def get_record_id(r):
        # This record should be unique. Only one operation per volume is possible
        #TODO: verify actual DB duplicate constraint
        return "%s-%s" % (r['start'], r['volume'])
        # More complete record?
        # return "%s-%s-%s-%s-%s" % (r['start'], r['node'], r['type'], r['volume'], r['storage_group'])

    def get_records(self, limit=None):
        """Select the mounting records from the tape_mounts table

Database table:
::
    accounting=> \d tape_mounts
                   Table "public.tape_mounts"
        Column     |            Type             | Modifiers
    ---------------+-----------------------------+-----------
     node          | character varying           | not null
     volume        | character varying           | not null
     type          | character varying(32)       | not null
     logname       | character varying(16)       | not null
     start         | timestamp without time zone | not null
     finish        | timestamp without time zone |
     state         | character(1)                | not null
     storage_group | character varying           |
     reads         | integer                     |
     writes        | integer                     |
    Indexes:
        "tape_mnts_node_idx" btree (node)
        "tape_mnts_oid_idx" btree (oid)
        "tape_mnts_start_idx" btree (start)
        "tape_mnts_type_idx" btree (type)
        "tape_mnts_volume_idx" btree (volume)
        "tape_mounts_logname_idx" btree (logname)
        "tape_mounts_storage_group_idx" btree (storage_group)
        The time zone of the timestamps is the local one

Records are selected and transformed for the probe consumption.

:param limit: maximum number of hours to include in the query
:yield: record for the probe
        """
        #OK_TO_SEND_RECORD = {'state': 'oktosend', 'volume': 'oktosend'}
        checkpoint = self.checkpoint
        where_clauses = []
        # DB uses local time -> checkpoint and all timestamps are in local time
        start_time = None
        # initialized anyway below -  end_time = datetime.datetime.now()
        ok_to_send_time = None
        end_by_limit = False

        if checkpoint:
            start_time = checkpoint.date()
            retrieved_previous_end_time = checkpoint.aux()
            checkpoint_interval = checkpoint.transaction()
            if checkpoint_interval is None:
                checkpoint_interval = 0
            else:
                checkpoint_interval = int(checkpoint_interval)
            if checkpoint_interval == 0:
                # NO interval, OK to send from the beginning
                DebugPrint(4, "Sending OK_TO_SEND - no interval in checkpoint (%s/%s)" %
                           (start_time, checkpoint_interval))
                yield EnstoreTapeDriveInput.OK_TO_SEND_RECORD
            else:
                estimated_previous_end_time = timeutil.wind_time(start_time, seconds=checkpoint_interval,
                                                                 backward=False)
                if retrieved_previous_end_time is not None and \
                        retrieved_previous_end_time != estimated_previous_end_time:
                    DebugPrint(4, "End_time in checkpoint not matching: estimated:%s, retrieved:%s" %
                               (estimated_previous_end_time, retrieved_previous_end_time))
                ok_to_send_time = estimated_previous_end_time
            DebugPrint(4, "Loaded checkpoint: %s (-%s), %s (%s - %s)" %
                       (start_time, self.rollback, ok_to_send_time, checkpoint_interval, retrieved_previous_end_time))
            if self.rollback > 0:
                start_time = timeutil.wind_time(start_time, seconds=self.rollback)
            where_clauses.append("start >= '%s'" % timeutil.format_datetime(start_time, iso8601=False))
        else:
            # NO Checkpoint - OK to send from the beginning
            DebugPrint(4, "Sending OK_TO_SEND - no checkpoint")
            yield EnstoreTapeDriveInput.OK_TO_SEND_RECORD

        if limit > 0:
            end_time = timeutil.wind_time(start_time, hours=limit, backward=False)
            # If input_delay is 0, check that the end_time is not in the future
            delay_time = timeutil.wind_time(datetime.datetime.now(), seconds=self.input_delay)
            if end_time > delay_time:
                end_time = delay_time
            else:
                end_by_limit = True
            #end_time = min(end_time, timeutil.wind_time(datetime.datetime.now(), seconds=self.input_delay))
        else:
            end_time = timeutil.wind_time(datetime.datetime.now(), seconds=self.input_delay)
        if checkpoint or limit or self.input_delay>0:
            end_time = timeutil.at_minute(end_time)
            where_clauses.append("start < '%s'" % timeutil.format_datetime(end_time, iso8601=False))
            if ok_to_send_time is not None and ok_to_send_time >= timeutil.wind_time(end_time, seconds=60):
                # If ok_to_send_time is not None, then checkpoint_interval was assigned (in checkpoint block)
                DebugPrint(2, "End time comes before new records are encountered (%s - %s <= 60sec)" %
                           (end_time, ok_to_send_time))
                if end_by_limit:
                    DebugPrint(2, "To avoid misinterpreting records DataLengthMax in the config file must be > %s" %
                               (checkpoint_interval/3600))
                else:
                    DebugPrint(2, "Either the probe runs too frequently or DataLengthMax may be too short. "
                                  "Current interval (hours):" %
                               (checkpoint_interval/3600))
        if start_time is not None:
            if self.input_min_interval > 0:
                if start_time > timeutil.wind_time(end_time, seconds=self.input_min_interval):
                    return
            else:
                if start_time >= end_time:
                    return
        if where_clauses:
            where_sql = "WHERE %s" % " AND ".join(where_clauses)
        else:
            where_sql = ""

        sql = '''SELECT
            node,
            volume,
            type,
            logname,
            start,
            finish,
            state,
            storage_group,
            reads,
            writes
            FROM tape_mounts
            %s
            ORDER BY start, storage_group
            ''' % (where_sql, )

        DebugPrint(4, "Requesting new EnstoreTapeDrive records %s" % sql)
        last_record_start_time = None
        first_record_start_time = None
        first_record = None
        mount_checkpoint = {}
        for r in self.query(sql):
            # Filter out values that are not acceptable
            if r['storage_group'] is None:
                continue
            if r['state'] not in ('M', 'D'):
                continue
            if ok_to_send_time is not None:
                # if ok_to_send_time is not None, checkpoint is True
                if r['start'] >= ok_to_send_time:
                    # send also the current record (yield is after)
                    # Time intervals are closed on the left (start) and open on the right (end)
                    yield EnstoreTapeDriveInput.OK_TO_SEND_RECORD
                    # to send the record OK_TO_SEND_RECORD only once
                    ok_to_send_time = None
            yield r
            if checkpoint:
                state = r['state']
                last_record_start_time = r['start']  # using start because finish could be NULL
                if first_record is None:
                    first_record = r
                    first_record_start_time = last_record_start_time
                if state == 'M':
                    mount_checkpoint[r['volume']] = last_record_start_time
                elif state == 'D':
                    mount_checkpoint[r['volume']] = None

        if last_record_start_time is not None and end_time is not None:
            # Looking 6mo before 4.27.2015 there are an average of 167 M od R records per hour
            if timeutil.wind_time(end_time, minutes=10) > last_record_start_time:
                DebugPrint(3, "Warning, no records in the last 10 min of EnstoreTapeDrive probe (%s - %s)" %
                           (end_time, last_record_start_time))

        if checkpoint:
            # first_unresolved_mount if any should be before end_time
            first_unresolved_mount = end_time
            for i in mount_checkpoint.values():
                if i is not None and i < first_unresolved_mount:
                    # If there are records (exist a i not None), first record variables are guaranteed not None
                    if i > first_record_start_time or not end_by_limit:
                        # it is past the first record or the time span is shorter than DataLengthMax
                        # this guarantees that the new invocation will may have more records
                        first_unresolved_mount = i
                    else:
                        # skip i for checkpoint consideration
                        DebugPrint(2,
                                   "Warning, reached DataLengthMax while the first mount record is still not matched.\n"
                                   "Sending oktosend, a fake dismount record and advancing the checkpoint past it")
                        yield EnstoreTapeDriveInput.OK_TO_SEND_RECORD
                        yield self.get_dismount_mount_record(first_record, end_time)

            if first_unresolved_mount == end_time:
                checkpoint_interval = 0
            else:
                checkpoint_interval = timeutil.total_seconds(end_time-first_unresolved_mount)
            DebugPrint(4, "Saving new EnstoreTapeDrive checkpoint: %s - %s (%s)" %
                       (first_unresolved_mount, end_time, checkpoint_interval))
            checkpoint.set_date_transaction_aux(first_unresolved_mount, checkpoint_interval, end_time)

    def _get_records_stub(self, limit=None):
        """get_records replacement for tests: records are from a pre-filled array
        limit is ignored"""
        DebugPrint(4, "Stub function: ignoring checkpoint and limit")
        yield EnstoreTapeDriveInput.OK_TO_SEND_RECORD
        for i in _EnstoreTapeDriveInputStub.get_records():
            yield i
        # To test checkpoint saving
        checkpoint = self.checkpoint
        if checkpoint:
            DebugPrint(4, "Stub function, saving checkpoint with same date: %s" % checkpoint.date())
            checkpoint.set_date_transaction_aux(checkpoint.date())

    def do_test(self, static_info=None):
        """Test with pre-arranged DB query results
        replacing: start, get_records
        """
        # replace DB calls with stubs
        self.start = self._start_stub
        self.get_records = self._get_records_stub


class EnstoreTapeDriveProbe(GratiaMeter):

    PROBE_NAME = 'enstore-tapedrive'
    # dCache, xrootd, Enstore
    SE_NAME = 'Enstore'
    # Production
    SE_STATUS = 'Production'
    # disk, tape
    SE_TYPE = 'tape'
    # raw, logical
    SE_MEASUREMENT_TYPE = 'logical'

    def __init__(self):
        GratiaMeter.__init__(self, self.PROBE_NAME)
        self._probeinput = EnstoreTapeDriveInput()

    @staticmethod
    def mount_complete_and_send(mrecord, finish, reads=0, writes=0, estimated=False, max_duration=0, min_duration=0):
        """Record the mount operation (estimated if the mount record is missing)

        :param mrecord: mount or srecord object (dictionary) - only common keys are used -
        :param finish: finish of the dismount operation (real or estimated), datetime
        :param estimated: 'estimated' if mount or dismount record were missing, 'ok' otherwise
        :return: Gratia record

        mount record:
            'type': srecord['type'],
            'volume': srecord['volume'],
            'node': srecord['node'],
            'storage_group': srecord['storage_group'],
            'mount_start': srecord['start'],
            'mount_finish': srecord['finish'],
            'estimated': estimated

        """

        DebugPrint(5, "Processing tape drive mount record: %s, FINISH: %s" % (mrecord, finish))

        # Empty usage record
        r = Gratia.UsageRecord("TapeDrive")
        r.Grid("Local")

        # The record Must have LocalUserId otherwise is quarantined. Adding a fake one
        r.LocalUserId('enstore')

        r.VOName(mrecord['storage_group'])

        # Naive timestamps (datetime obj) with actual local time zone (assume_local=True is the default)
        start = timeutil.datetime_to_unix_time(timeutil.datetime_to_utc(mrecord['mount_start'], naive=True))
        if not finish:
            finish = start
            duration = 0
        else:
            finish = timeutil.datetime_to_unix_time(timeutil.datetime_to_utc(finish, naive=True))
            duration = int(float(finish)-float(start))


        # Adding ID. Here, so finish is defined and is the # of seconds from epoch
        # TODO: is the ID OK? Unique enough?
        local_id = "%s-%s-%s-%s" % (mrecord['node'], mrecord['type'], mrecord['volume'], finish)

        # Status is 'estimated' if either the mount or dismount record are missing, 'ok' otherwise
        # estimated status means that the duration is estimated
        # Status in UR is varchar 255
        if estimated or mrecord['estimated']:
            r.Status('estimated')
            # calculating floor/ceiling limits for the estimated duration
            if max_duration > 0 and duration > max_duration:
                DebugPrint(3, "Capping mount record (%s) from %s to %s" % (local_id, duration, max_duration))
                duration = max_duration
                finish = start + duration
                # finish changed, updating the ID
                local_id = "%s-%s-%s-%s" % (mrecord['node'], mrecord['type'], mrecord['volume'], finish)
            elif duration < min_duration:
                DebugPrint(3, "Increasing min length of mount record (%s) from %s to %s" %
                           (local_id, duration, max_duration))
                duration = min_duration
                start = finish - duration
        else:
            r.Status('ok')
        r.WallDuration(duration)
        r.StartTime(timeutil.format_datetime(start))
        r.EndTime(timeutil.format_datetime(finish))

        r.LocalJobId(local_id)
        r.GlobalJobId(local_id)

        r.AdditionalInfo("reads", reads)
        r.AdditionalInfo("writes", writes)

        r.SubmitHost(mrecord['node'])
        r.Queue(mrecord['volume'])

        # Future modifications of Enstore may include a DN
        # r.DN("/OU=UnixUser/CN=%s" % srecord['username'])

        DebugPrint(4, "Sending tape drive record for VO %s: %s" % (mrecord['storage_group'], local_id))
        Gratia.Send(r)

    @staticmethod
    def mount_init(srecord, estimated=False):
        """Create a mount record copying the values form the input

        :param srecord: mount or dismount record (dictionary)
        :param estimated: 'estimated' if mount or dismount record were missing, 'ok' otherwise
        :return: mount record
        """
        retv = {
            'type': srecord['type'],
            'volume': srecord['volume'],
            'node': srecord['node'],
            'storage_group': srecord['storage_group'],
            'mount_start': srecord['start'],
            'mount_finish': srecord['finish'],
            'estimated': estimated
        }
        return retv

    def main(self):
        """Main loop: retrieve the data and send it to the Gratia collector"""
        # Initialize the probe an the input
        self.start()
        DebugPrint(4, "Enstore tape drive probe started")

        #se = self.get_sitename()
        #name = self.get_probename()
        #hostname = self.get_hostname()

        # max and min estimated mount time. The limits are not used in case of real mount time.
        max_mount_time = int(self.get_config_attribute('EnstoreMaxMountTime', 0))
        min_mount_time = int(self.get_config_attribute('EnstoreMinMountTime', 0))


        # Variable used to unlock sending of records.
        # Controlled by the input via special srecord with status='oktosend'
        ok_to_send = False
        mounts = {}

        # Loop over storage records
        for srecord in self._probeinput.get_records(self.get_input_max_length()):
            """Values in srecord:
            node,
            volume, - label
            type, -
            logname,
            start,
            finish,
            state,
            storage_group,
            reads,
            writes
            """
            DebugPrint(5, "Preparing tape drive record for: %s" % srecord)

            # volume (or label) is a unique identifier of a tape.
            # a tape can be mounted (only dismount is possible) or unmounted (only mount is possible)
            # it cannot be mounted more than once by different storage_groups
            #  mounts[t_volume][t_type, storage_group, ...]
            #t_type = srecord['type']
            #storage_group = srecord['storage_group']
            t_volume = srecord['volume']

            if srecord['state'] == 'M':
                # mount of the tape
                if t_volume in mounts:
                    # An already mounted tape is mounted. A dismount record may have been lost
                    DebugPrint(3, "Tape %s already mounted by %s at %s. Mount operation by %s reported at %s" % (
                        t_volume, mounts[t_volume]['storage_group'], mounts[t_volume]['mount_start'],
                        srecord['storage_group'], srecord['start'])
                    )
                    # Send the old record (end at the end of mount OR at start of new mount) and remove the mount
                    # "if you loose the ticket you pay all OR get by"
                    if ok_to_send:
                        EnstoreTapeDriveProbe.mount_complete_and_send(mounts[t_volume], srecord['start'],
                                                                      estimated=True, max_duration=max_mount_time,
                                                                      min_duration=min_mount_time)
                    else:
                        DebugPrint(4, "Not sending duplicate record: %s, %s" % (mounts[t_volume], srecord['start']))
                    # del not needed here because is reassigned right after anyway
                    # del mounts[t_volume]
                # update the record
                mounts[t_volume] = EnstoreTapeDriveProbe.mount_init(srecord)
                # read and write information is reported only at dismount
            elif srecord['state'] == 'D':
                # dismount of the tape
                try:
                    if mounts[t_volume]['storage_group'] != srecord['storage_group']:
                        DebugPrint(3, "Storage Group (VO) of tape %s not consistent: %s at mount, %s at dismount." % (
                            t_volume, mounts[t_volume]['storage_group'], srecord['storage_group'])
                        )
                        # close, dismount at finish, send the record, and remove the mount
                        if ok_to_send:
                            EnstoreTapeDriveProbe.mount_complete_and_send(mounts[t_volume], srecord['start'],
                                                                          estimated=True, max_duration=max_mount_time,
                                                                          min_duration=min_mount_time)
                        else:
                            DebugPrint(4, "Not sending duplicate record: %s, %s" % (mounts[t_volume], srecord['start']))
                        del mounts[t_volume]
                except KeyError:
                    pass
                if not t_volume in mounts:
                    # A mount record was missed
                    DebugPrint(3, "Tape %s not mounted. Dismount operation by %s reported at %s" % (
                        t_volume, srecord['storage_group'], srecord['start'])
                    )
                    # assuming a mount record starting (and finishing) at the beginning of dismount
                    mounts[t_volume] = EnstoreTapeDriveProbe.mount_init(srecord, True)
                if mounts[t_volume]['type'] != srecord['type']:
                    DebugPrint(3, "Type of tape %s not consistent: %s at mount, %s at dismount." % (
                        t_volume, mounts[t_volume]['type'], srecord['type'])
                    )
                try:
                    mount_reads = int(srecord['reads'])
                except (ValueError, TypeError):
                    mount_reads = 0
                try:
                    mount_writes = int(srecord['writes'])
                except (ValueError, TypeError):
                    mount_writes = 0
                #tmp_record = mounts.pop(t_volume) # instead of access and del
                if ok_to_send:
                    EnstoreTapeDriveProbe.mount_complete_and_send(mounts[t_volume], srecord['finish'],
                                                                  mount_reads, mount_writes,
                                                                  max_duration=max_mount_time,
                                                                  min_duration=min_mount_time)
                else:
                    DebugPrint(5, "Not sending duplicate record: %s, %s" % (mounts[t_volume], srecord['start']))
                # delete anyway
                del mounts[t_volume]
            elif srecord['state'] == 'oktosend':
                DebugPrint(4, "Received OK_TO_SEND - setting True")
                ok_to_send = True
            else:
                # Invalid state!
                DebugPrint(3, "Invalid state %s. Operation on volume %s by %s reported at %s" % (
                    srecord['state'], t_volume,
                    srecord['storage_group'], srecord['start'])
                )


if __name__ == "__main__":
    # Do the work
    EnstoreTapeDriveProbe().main()



