###############################################################################
# (c) Copyright 2019 CERN for the benefit of the LHCb Collaboration           #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "LICENSE".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
"""RAWIntegrityDB is the front end for the database containing the files which
are awating migration.

It offers a simple interface to add files, get files and modify their
status.
"""
from DIRAC import gLogger, S_OK, S_ERROR
from DIRAC.Core.Base.DB import DB
from DIRAC.Core.Utilities.List import stringListToString

gLogger.initialize("DMS", "/Databases/RAWIntegrityDB/Test")


class RAWIntegrityDB(DB):
    _tablesDict = {}
    # Files table
    _tablesDict["Files"] = {
        "Fields": {
            "ID": "BIGINT NOT NULL AUTO_INCREMENT",
            "LFN": "VARCHAR(255) NOT NULL",
            "PFN": "VARCHAR(255) NOT NULL",
            "Size": "BIGINT NOT NULL",
            "StorageElement": "VARCHAR(32) NOT NULL",
            "GUID": "VARCHAR(255) NOT NULL",
            "FileChecksum": "VARCHAR(255) NOT NULL",
            "SubmitTime": "DATETIME NOT NULL",
            "CompleteTime": "DATETIME",
            "Status": 'VARCHAR(255) DEFAULT "Active"',
        },
        "PrimaryKey": "ID",
        "Indexes": {"Status": ["Status"], "LFN": ["LFN"]},
    }
    # LastMonitor table
    _tablesDict["LastMonitor"] = {"Fields": {"LastMonitorTime": "DATETIME NOT NULL PRIMARY KEY"}}

    def __init__(self, systemInstance="Default"):
        DB.__init__(self, "RAWIntegrityDB", "DataManagement/RAWIntegrityDB")

    def _checkTable(self):
        """_checkTable.

        Make sure the table is created
        """

        return self.__createTables()

    def __createTables(self):
        """__createTables.

        Writes the schema in the database. If a table is already in the
        schema, it is skipped to avoid problems trying to create a table
        that already exists.
        """

        # Horrible SQL here !!
        existingTables = self._query("show tables")
        if not existingTables["OK"]:
            return existingTables
        existingTables = [existingTable[0] for existingTable in existingTables["Value"]]

        # Makes a copy of the dictionary _tablesDict
        tables = {}
        tables.update(self._tablesDict)

        for existingTable in existingTables:
            if existingTable in tables:
                del tables[existingTable]

        res = self._createTables(tables)
        if not res["OK"]:
            return res

        # Human readable S_OK message
        if res["Value"] == 0:
            res["Value"] = "No tables created"
        else:
            res["Value"] = f"Tables created: {','.join(list(tables))}"

            # Create the initial value
            if "LastMonitor" in tables:
                ret = self._query("INSERT INTO LastMonitor(LastMonitorTime) value (UTC_TIMESTAMP());")
                if not ret["OK"]:
                    return ret

        return res

    def showTables(self):
        """return the list of tables."""

        existingTables = self._query("show tables")
        if not existingTables["OK"]:
            return existingTables
        existingTables = [existingTable[0] for existingTable in existingTables["Value"]]
        return S_OK(existingTables)

    def getActiveFiles(self):
        """Obtain all the active files in the database along with all their
        associated metadata."""
        try:
            gLogger.info("RAWIntegrityDB.getActiveFiles: " "Obtaining files awaiting migration from database.")
            res = self._query(
                "SELECT "
                "LFN,PFN,Size,StorageElement,GUID,FileChecksum,TIME_TO_SEC(TIMEDIFF(UTC_TIMESTAMP(),SubmitTime)) "
                "from Files WHERE Status = 'Active';"
            )
            if not res["OK"]:
                gLogger.error("RAWIntegrityDB.getActiveFiles: Failed to get files from database.", res["Message"])
                return res
            else:
                fileDict = {}
                for lfn, pfn, size, se, guid, checksum, waittime in res["Value"]:
                    fileDict[lfn] = {
                        "PFN": pfn,
                        "Size": size,
                        "SE": se,
                        "GUID": guid,
                        "Checksum": checksum,
                        "WaitTime": waittime,
                    }
                gLogger.info(
                    f"RAWIntegrityDB.getActiveFiles: Obtained {len(fileDict)} files awaiting migration from database."
                )
                return S_OK(fileDict)
        except Exception as x:
            errStr = "RAWIntegrityDB.getActiveFiles: Exception while getting files from database."
            gLogger.exception(errStr, lException=x)
            return S_ERROR(errStr)

    def getFiles(self, status):
        """Obtain all the files in a give status (or list) in the database along
        with all their associated metadata.

        :param status: string or list of status we want to query

        :returns: dict {lfn: metadata}
        """
        if isinstance(status, str):
            status = [status]

        try:
            gLogger.info(f"RAWIntegrityDB.getActiveFiles: Obtaining files in status {status}.")
            res = self._query(
                "SELECT "
                "LFN,Status,PFN,Size,StorageElement,GUID,FileChecksum,"
                "TIME_TO_SEC(TIMEDIFF(UTC_TIMESTAMP(),SubmitTime)),SubmitTime "
                "from Files WHERE Status in (%s);" % (stringListToString(status))
            )
            if not res["OK"]:
                gLogger.error("RAWIntegrityDB.getActiveFiles: Failed to get files from database.", res["Message"])
                return res
            fileDict = {}
            for lfn, st, pfn, size, se, guid, checksum, waittime, submitTime in res["Value"]:
                fileDict[lfn] = {
                    "PFN": pfn,
                    "Size": size,
                    "SE": se,
                    "GUID": guid,
                    "Checksum": checksum,
                    "WaitTime": waittime,
                    "Status": st,
                    "SubmitTime": submitTime,
                }
            gLogger.info(
                f"RAWIntegrityDB.getActiveFiles: Obtained {len(fileDict)} files awaiting migration from database."
            )
            return S_OK(fileDict)
        except Exception as x:
            errStr = "RAWIntegrityDB.getActiveFiles: Exception while getting files from database."
            gLogger.exception(errStr, lException=x)
            return S_ERROR(errStr)

    def getUnmigratedFiles(self):
        """Get files in status Active, Copied, and Registered."""

        return self.getFiles(["Active", "Copied", "Registered"])

    def setFileStatus(self, lfn, status):
        """Update the status of the file in the database.

        A file in status 'Done' cannot be unchanged
        """
        try:
            gLogger.info(f"RAWIntegrityDB.setFileStatus: Attempting to update status of {lfn} to '{status}'.")
            req = "UPDATE Files SET Status='{}',CompleteTime=UTC_TIMESTAMP() WHERE LFN = '{}' AND Status != 'Done';".format(
                status,
                lfn,
            )
            res = self._update(req)
            if not res["OK"]:
                gLogger.error("RAWIntegrityDB.setFileStatus: Failed update file status.", res["Message"])
            else:
                gLogger.info("RAWIntegrityDB.setFileStatus: Successfully updated file status.")
            return res
        except Exception as x:
            errStr = "RAWIntegrityDB.setFileStatus: Exception while updating file status."
            gLogger.exception(errStr, lException=x)
            return S_ERROR(errStr)

    def removeFile(self, lfn):
        """Remove file from the DB."""
        try:
            gLogger.info(f"RAWIntegrityDB.removeFile: Attempting to remove {lfn}.")
            req = f"DELETE FROM Files WHERE LFN = '{lfn}';"
            res = self._update(req)
            if not res["OK"]:
                gLogger.error("RAWIntegrityDB.removeFile: Failed to remove file.", res["Message"])
            else:
                gLogger.info("RAWIntegrityDB.removeFile: Successfully removed file.")
            return res
        except Exception as x:
            errStr = "RAWIntegrityDB.removeFile: Exception while removing file."
            gLogger.exception(errStr, lException=x)
            return S_ERROR(errStr)

    def addFile(self, lfn, pfn, size, se, guid, checksum):
        """Insert file into the database."""
        try:
            gLogger.info(f"RAWIntegrityDB.addFile: Attempting to add {lfn} to database.")
            req = (
                "INSERT INTO Files (LFN,PFN,Size,StorageElement,GUID,FileChecksum,SubmitTime) VALUES\
            ('%s','%s',%s,'%s','%s','%s',UTC_TIMESTAMP());"
                % (lfn, pfn, size, se, guid, checksum)
            )
            res = self._update(req)
            if not res["OK"]:
                gLogger.error("RAWIntegrityDB.addFile: Failed update add file to database.", res["Message"])
            else:
                gLogger.info("RAWIntegrityDB.addFile: Successfully added file.")
            return res
        except Exception as x:
            errStr = "RAWIntegrityDB.addFile: Exception while updating file status."
            gLogger.exception(errStr, lException=x)
            return S_ERROR(errStr)

    def setLastMonitorTime(self):
        """Set the last time the migration rate was calculated."""
        try:
            gLogger.info("RAWIntegrityDB.setLastMonitorTime: Attempting to set the last migration marker.")
            req = "UPDATE LastMonitor SET LastMonitorTime=UTC_TIMESTAMP();"
            res = self._update(req)
            if not res["OK"]:
                gLogger.error("RAWIntegrityDB.setLastMonitorTime: Failed update migration marker.", res["Message"])
            else:
                gLogger.info("RAWIntegrityDB.setLastMonitorTime: Successfully updated migration marker.")
            return res
        except Exception as x:
            errStr = "RAWIntegrityDB.setLastMonitorTime: Exception while updating migration marker."
            gLogger.exception(errStr, lException=x)
            return S_ERROR(errStr)

    def getLastMonitorTimeDiff(self):
        """Get the last time the migration rate was calculated."""
        try:
            gLogger.info("RAWIntegrityDB.getLastMonitorTimeDiff: Attempting to get the last migration marker.")
            req = "SELECT TIME_TO_SEC(TIMEDIFF(UTC_TIMESTAMP(),LastMonitorTime)) FROM LastMonitor LIMIT 1;"
            res = self._query(req)
            if not res["OK"]:
                gLogger.error("RAWIntegrityDB.getLastMonitorTimeDiff: Failed get migration marker.", res["Message"])
                return res
            else:
                gLogger.info("RAWIntegrityDB.getLastMonitorTimeDiff: Successfully obtained migration marker.")
                timediff = res["Value"][0][0]
                return S_OK(timediff)
        except Exception as x:
            errStr = "RAWIntegrityDB.getLastMonitorTimeDiff: Exception while getting migration marker."
            gLogger.exception(errStr, lException=x)
            return S_ERROR(errStr)

    def getGlobalStatistics(self):
        """Get the count of the file statutes in the DB."""
        req = "SELECT Status,COUNT(*) FROM Files GROUP BY Status;"
        res = self._query(req)
        if not res["OK"]:
            return res
        statusDict = {}
        for resValue in res["Value"]:
            status, count = resValue
            statusDict[status] = count
        return S_OK(statusDict)

    def getFileSelections(self):
        """Get the unique values of the selection fields."""
        selDict = {"StorageElement": [], "Status": []}
        req = "SELECT DISTINCT(StorageElement) FROM Files;"
        res = self._query(req)
        if not res["OK"]:
            return res
        for resValue in res["Value"]:
            selDict["StorageElement"].append(resValue[0])
        req = "SELECT DISTINCT(Status) FROM Files;"
        res = self._query(req)
        if not res["OK"]:
            return res
        for resValue in res["Value"]:
            selDict["Status"].append(resValue[0])
        return S_OK(selDict)

    def __buildCondition(self, condDict, older=None, newer=None, timeStamp="SubmitTime"):
        """build SQL condition statement from provided condDict and other extra
        conditions."""
        condition = ""
        conjunction = "WHERE"
        if condDict is not None:
            for attrName, attrValue in condDict.items():
                ret = self._escapeString(attrName)
                if not ret["OK"]:
                    return ret
                attrName = "`" + ret["Value"][1:-1] + "`"
                if isinstance(attrValue, list):
                    multiValueList = []
                    for x in attrValue:
                        ret = self._escapeString(x)
                        if not ret["OK"]:
                            return ret
                        x = ret["Value"]
                        multiValueList.append(x)
                    multiValue = ",".join(multiValueList)
                    condition = f" {condition} {conjunction} {attrName} in ({multiValue})"
                else:
                    ret = self._escapeString(attrValue)
                    if not ret["OK"]:
                        return ret
                    attrValue = ret["Value"]
                    condition = f" {condition} {conjunction} {attrName}={attrValue}"
                conjunction = "AND"
        if older:
            ret = self._escapeString(older)
            if not ret["OK"]:
                return ret
            older = ret["Value"]
            condition = f" {condition} {conjunction} {timeStamp} <= {older}"
            conjunction = "AND"
        if newer:
            ret = self._escapeString(newer)
            if not ret["OK"]:
                return ret
            newer = ret["Value"]
            condition = f" {condition} {conjunction} {timeStamp} >= {newer}"
        return condition

    def selectFiles(self, selectDict, orderAttribute="LFN", newer=None, older=None, limit=None):
        """Select the files which match the selection criteria."""
        condition = self.__buildCondition(selectDict, older, newer)
        if orderAttribute:
            orderType = None
            orderField = orderAttribute
            if orderAttribute.find(":") != -1:
                orderType = orderAttribute.split(":")[1].upper()
                orderField = orderAttribute.split(":")[0]
            condition = condition + " ORDER BY " + orderField
            if orderType:
                condition = condition + " " + orderType
        if limit:
            condition = condition + " LIMIT " + str(limit)
        res = self._query(
            "SELECT "
            "LFN,PFN,Size,StorageElement,GUID,FileChecksum,SubmitTime,CompleteTime,Status "
            "from Files %s" % condition
        )
        if not res["OK"]:
            return res
        if not len(res["Value"]):
            return S_OK([])
        return S_OK(res["Value"])


# ...............................................................................
