# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Queue Manager.

This is a slight misnomer: this code really only manages taking things
out of the queue.  But I can't think of a better name right now.

It is not in the same module as the code that puts things into the
queue, as these are rarely both used in the same program.

"""

import os
import time
import shutil
import filecmp
import logging

from sigma.util import avalog
from sigma.spool.spooling import Queue

logger = logging.getLogger("spool.qmgr")

class QueueManager(Queue):

    def __init__(self, spooldir=None, options=None):
        """Constructor."""
        Queue.__init__(self, spooldir, options)

    def is_empty(self):
        """Returns True iff the queue is empty."""
        return not self.is_non_empty()

    def is_non_empty(self):
        """Returns True iff the queue is NOT (!) empty."""
        names = os.listdir(self.spooldir)
        names = list(map(os.path.normcase, names))
        nameset = dict.fromkeys(names)
        for name in names:
            if self.make_item(name, nameset) is not None:
                return True
        return False

    def count_items(self):
        """Counts the number of queue items.

        The result may be off if broken queue items exist (X.spool for
        which X doesn't exist or where the contents of X.spool is
        invalid); it's too expensive to validate every item for the
        purpose for which this API is typically called.  If you simply
        want to know whether the queue is empty or not, use is_empty()
        or is_non_empty() instead; those *do* validate at least one
        item.

        """
        names = os.listdir(self.spooldir)
        names = list(map(os.path.normcase, names))
        count = 0
        for name in names:
            if name.endswith(".spool"):
                count += 1
        return count

    def age_stats(self):
        """Computes age statistics for all queue items.

        This returns a tuple (count, oldest, average, youngest) where:

        - count is the number of spool queue items
        - oldest is the mtime of the oldest item
        - average is the average mtime of all items
        - youngest is the mtime of the youngest item

        If count is zero, the remaining values are also zero.

        The same caveat as for count_items() applies: the result may
        be off if broken queue items exist.

        If you don't need the mtimes, count_items() is more efficient,
        since it doesn't stat each file.

        """
        names = os.listdir(self.spooldir)
        names = list(map(os.path.normcase, names))
        count = 0
        oldest = None
        youngest = 0
        total = 0
        for name in names:
            if name.endswith(".spool"):
                fullname = os.path.join(self.spooldir, name)
                try:
                    mtime = os.path.getmtime(fullname)
                except os.error:
                    continue # Skip unstattable items
                count += 1
                if oldest is None or mtime < oldest:
                    oldest = mtime
                if mtime > youngest:
                    youngest = mtime
                total += mtime
        if not count:
            return 0, 0, 0, 0
        average = float(total) / count
        return count, oldest, average, youngest

    def list_items(self):
        """A generator yielding successive queue items.

        This only yields queue items whose filename exists and whose
        metadata is consistent and indicates that the item is ready
        for upload.

        """
        names = os.listdir(self.spooldir)
        names = list(map(os.path.normcase, names))
        names.sort()
        nameset = dict.fromkeys(names)
        for name in names:
            qitem = self.make_item(name, nameset)
            if qitem is not None:
                yield qitem

    def make_item(self, name, nameset):
        """Helper to produce a queue item from a spool filename.

        The argument must be a normalized file name (no path info).

        Returns a ManagedQueueItem or None if the queue item is
        invalid.

        """
        if name.endswith(".spool"):
            qitem = ManagedQueueItem(os.path.join(self.spooldir, name))
            qfile = os.path.normcase(qitem.metadata.get("qfile", ""))
            if (qitem.status == "queued" and
                qfile in nameset and
                qfile == name[:-len(".spool")]):
                    return qitem
        return None

    def full_clear(self):
        """Completely cleans out the queue.

        This removes all files and directories from the spool
        area, regardless of age.  Use with discretion!

        For a gentler cleanup, see periodic_cleanup().

        """
        for name in os.listdir(self.spooldir):
            filename = os.path.join(self.spooldir, name)
            if os.path.isdir(filename):
                logger.info("Removing directory tree %r", filename)
                shutil.rmtree(filename, False, self.onerror)
            else:
                try:
                    os.remove(filename)
                except os.error as msg:
                    logger.warning("Can't remove %r: %s", filename, msg)
                else:
                    logger.info("Discarded %r", filename)

    def onerror(self, func, pathname, exc_tuple):
        logger.warning("Can't %s %s: %s",
                       func.__name__, pathname, exc_tuple[1])

    def periodic_cleanup(self, N1=5*60, N2=8*3600, N3=4*24*3600):
        """Does periodic cleanup.

        Delete files from the spool queue that are left behind by
        crashed processes, or that simply don't belong there.

        There are three parameters, N1, N2, N2, which give timeouts in
        seconds.  We require N3 >= 10*N2 >= 100*N1.  The defaults are:

        - N1 = 5 minutes; used to delete garbage files
        - N2 = 8 hours; used to delete open spool items
        - N3 = 4 days; used to delete queued spool items

        Garbage files are:

        - files that are neither data nor spool files
        - data files without spool files
        - spool files without data files
        - data/spool pairs with damaged spool file

        Note that files whose extension is either .open and .spool are
        considered spool files for this purpose.

        Note that during the process of rewriting the spool file, the
        spool file may temporarily look damaged (it is truncated first
        and then rewritten); during the process of creating a spool
        item, a data file may exist without a spool file.  These
        conditions only last for a fraction of a second normally.

        See also limit_quota().

        """
        assert 10*N1 <= N2, (N1, N2)
        assert 10*N2 <= N3, (N2, N3)
        names = os.listdir(self.spooldir)
        names = list(map(os.path.normcase, names))
        names.sort()
        nameset = dict.fromkeys(names)
        spooldir = os.path.normcase(self.spooldir)
        spoolprefix = os.path.normcase("spool")
        spoolsuffix = os.path.normcase(".spool")
        opensuffix = os.path.normcase(".open")
        def remove(fn, reason):
            try:
                os.remove(fn)
            except os.error as msg:
                logger.warning("Can't remove %r (%s): %s", fn, reason, msg)
            else:
                logger.info("Removed %r (%s)", fn, reason)
        for name in names:
            filename = os.path.join(spooldir, name)
            if os.path.isdir(filename):
                logger.debug("Skipping directory %r", filename)
                continue
            try:
                mtime = os.path.getmtime(filename)
            except os.error:
                # File disappeared -- ah well, ignore it
                continue
            t0 = time.time()
            t1 = t0 - N1
            if mtime >= t1:
                logger.debug("Leaving young file %r alone", filename)
                continue
            if not name.startswith(spoolprefix):
                remove(filename, "bogus")
                continue
            if (not name.endswith(spoolsuffix) and
                not name.endswith(opensuffix)):
                # Data file
                if (name + spoolsuffix not in nameset and
                    name + opensuffix not in nameset):
                    remove(filename, "orphan data")
                # Else we'll deal with it as part of the data/spool item
                continue
            # Spool metadata file (open or queued)
            if name.endswith(spoolsuffix):
                dataname = name[:-len(spoolsuffix)]
                status = "queued"
            else:
                assert name.endswith(opensuffix)
                dataname = name[:-len(opensuffix)]
                status = "open"
            if dataname not in nameset:
                remove(filename, "orphan spool")
                continue
            datafilename = os.path.join(spooldir, dataname)
            try:
                datamtime = os.path.getmtime(datafilename)
            except os.error:
                # File disappeared (maybe it got uploaded?); ignore
                continue
            if status == "queued":
                t3 = t0 - N3
                if mtime >= t3 or datamtime >= t3:
                    logger.debug("Leaving queued item %r alone", filename)
                    continue
            elif status == "open":
                t2 = t0 - N2
                if mtime >= t2 or datamtime >= t2:
                    logger.debug("Leaving open item %r alone", filename)
                    continue
            remove(filename, "expired %s item" % status)
            remove(datafilename, "expired %s data" % status)

    def limit_quota(self, limit):
        """Throws away old files until the total is <= limit Kbytes.

        The limit is expressed in kilobytes.  File sizes are rounded
        up to kilobytes (1024).  For best results, first run
        periodic_cleanup().

        """
        kbyte = 1024
        items = []
        totalsize = 0
        for qitem in self.list_items():
            try:
                datasize = os.path.getsize(qitem.filename)
                metasize = os.path.getsize(qitem.metadatafilename)
                mtime = os.path.getmtime(qitem.metadatafilename)
            except os.error:
                # The item disappeared (perhaps it got uploaded); ignore
                continue
            itemsize = (datasize+kbyte-1)//kbyte + (metasize+kbyte-1)//kbyte
            totalsize += itemsize
            # Store the negative itemsize so when there are items of
            # identical age the largest item will be deleted first:
            # I'd rather delete one big item than a whole bunch of
            # small ones.
            items.append((mtime, -itemsize, qitem))
        if totalsize <= limit:
            logger.info("%d items: total size (%d) <= limit (%d) Kbytes",
                        len(items), totalsize, limit)
            return
        logger.info("%d items: total size (%d) > limit (%d) Kbytes",
                    len(items), totalsize, limit)
        items.sort()
        ndeleted = 0
        for mtime, negitemsize, qitem in items:
            logger.info("Discarding %s to reduce queue size", qitem.filename)
            qitem.discard()
            totalsize += negitemsize
            if totalsize <= limit:
                break

    def consolidate_duplicates(self):
        """Removes duplicate spool queue items.

        This scans the queue for consecutive identical items with the
        same type (and if the type is 'compliance', with the same
        deployment_id).  It will remove all items except for the last,
        and update its starttime metadata to be the starttime of the
        first deleted item.

        Special case:

        if an item with type 'flush' is present, all except the most
        recent item of each type is deleted.  This is used by the
        server when is too overloaded to catch up in a reasonable
        time.

        """
        bytype = {}
        for qitem in self.list_items():
            type = qitem.metadata["type"]
            if type.endswith(".incr"):
                continue
            if type == "compliance":
                type += "." + qitem.metadata.get("deployment_id", "?")
            if type in bytype:
                bytype[type].append(qitem)
            else:
                bytype[type] = [qitem]
        flush = "flush" in bytype
        for type, items in bytype.items():
            assert len(items) >= 1
            items.sort() # Queue items sort by end time
            if flush:
                for item in items[:-1]:
                    logger.info("Flushing %s item %s", type,
                                os.path.basename(item.metadatafilename))
                    item.discard()
                continue
            prev = prevsize = None
            for item in items:
                try:
                    itemsize = os.path.getsize(item.filename)
                    if (itemsize == prevsize and
                        prev is not None and
                        filecmp.cmp(prev.filename, item.filename, False)):
                        item.metadata["starttime"] = (
                            prev.metadata.get("starttime") or
                            prev.metadata["created"])
                        if "endtime" not in item.metadata:
                            item.metadata["endtime"] = item.metadata["created"]
                        if item.rewrite_metadata():
                            logger.info(
                                "Consolidated duplicate %s item: %s => %s",
                                type,
                                os.path.basename(prev.metadatafilename),
                                os.path.basename(item.metadatafilename))
                            prev.discard()
                    prev = item
                    prevsize = itemsize
                except IOError as err:
                    logger.warn("IOError comparing spool queue items: %s", err)
                    prev = prevsize = None

class ManagedQueueItem(object):

    def __init__(self, metadatafilename):
        """Constructor.  Takes a metadata filename argument."""
        if not metadatafilename.endswith(".spool"):
            raise ValueError("metadata filename must end with .spool")
        self.metadatafilename = metadatafilename
        self.filename = metadatafilename[:-len(".spool")]
        self.read_metadata()
        self.status = self.metadata.get("status")
        self.sorttime = (self.metadata.get("endtime") or
                         self.metadata.get("created"))

    def __lt__(self, other):
        """Comparator.  Used for sorting a (sub)queue by time."""
        if not isinstance(other, ManagedQueueItem):
            return NotImplemented
        return self.sorttime < other.sorttime

    def read_metadata(self):
        """Reads metadata from a given filename.

        Store a dict mapping keywords to values as self.metadata.

        """
        self.metadata = d = {}
        try:
            f = open(self.metadatafilename, "r")
        except IOError:
            # Don't die on errors
            return
        try:
            for line in f:
                if ":" in line:
                    key, value = line.split(":", 1)
                    d[key] = value.strip()
        finally:
            f.close()

    def rewrite_metadata(self):
        """Writes the metadata back, after presumed modification."""
        try:
            f = open(self.metadatafilename, "w")
        except IOError as err:
            logger.error("Can't rewrite metadata for %s: %s",
                         self.metadatafilename, err)
            return False
        try:
            for key_value in self.metadata.items():
                f.write("%s: %s\n" % key_value)
        finally:
            f.close()
        return True

    def discard(self):
        """Discards the files associated with the spool item."""
        # XXX Shouldn't this log success as well?
        try:
            os.remove(self.metadatafilename)
        except os.error as err:
            logger.warning("Can't delete %r: %s", self.metadatafilename, err)
        try:
            os.remove(self.filename)
        except os.error as err:
            logger.warning("Can't delete %r: %s", self.filename, err)
