import os
import time
import gzip
from logging.handlers import TimedRotatingFileHandler


class CompressedTimedRotatingFileHandler(TimedRotatingFileHandler):
    # Opens the rotated log file, compresses it to .gz, then deletes the original
    def gzip_logs(self, file_path):
        gz_path = f"{file_path}.gz"
        try:
            with open(file_path, 'rb') as f_in, gzip.open(gz_path, 'wb') as f_out:
                f_out.writelines(f_in)
            os.remove(file_path)
        except Exception:
            # If gzip failed, clean up partial .gz file
            if os.path.exists(gz_path):
                os.remove(gz_path)
            raise
    
    # NOTE: copied code from TimedRotatingFileHandler getFilesToDelete()
    # override the method in TimedRotatingFileHandler because after gzipping, the pattern in the original function no longer works
    def getFilesToDelete(self):
        """
        Determine the files to delete when rolling over.
        More specific than the earlier method, which just used glob.glob().
        """
        dirName, baseName = os.path.split(self.baseFilename)
        fileNames = os.listdir(dirName)
        result = []
        if self.namer is None:
            prefix = baseName + '.'
            plen = len(prefix)
            for fileName in fileNames:
                orig_filename = fileName
                # Remove .gz extension if present for pattern matching
                fileName = fileName[:-3] if fileName.endswith('.gz') else fileName
                if fileName[:plen] == prefix:
                    suffix = fileName[plen:]
                    if self.extMatch.fullmatch(suffix):
                        result.append(os.path.join(dirName, orig_filename))
        else:
            for fileName in fileNames:
                orig_filename = fileName
                # Remove .gz extension if present for pattern matching
                fileName = fileName[:-3] if fileName.endswith('.gz') else fileName
                # Our files could be just about anything after custom naming,
                # but they should contain the datetime suffix.
                # Try to find the datetime suffix in the file name and verify
                # that the file name can be generated by this handler.
                m = self.extMatch.search(fileName)
                while m:
                    dfn = self.namer(self.baseFilename + "." + m[0])
                    if os.path.basename(dfn) == fileName:
                        result.append(os.path.join(dirName, orig_filename))
                        break
                    m = self.extMatch.search(fileName, m.start() + 1)

        if len(result) < self.backupCount:
            result = []
        else:
            result.sort()
            result = result[:len(result) - self.backupCount]
        return result
        
    # NOTE: copied code from TimedRotatingFileHandler doRollover()
    # used to get the dfn (destination filename) and pass it to gzip_logs()
    def doRollover(self):
        # get the time that this sequence started at and make it a TimeTuple
        currentTime = int(time.time())
        t = self.rolloverAt - self.interval
        if self.utc:
            timeTuple = time.gmtime(t)
        else:
            timeTuple = time.localtime(t)
            # dst = daylight saving time
            dstNow = time.localtime(currentTime)[-1]
            dstThen = timeTuple[-1]
            if dstNow != dstThen:
                if dstNow:
                    addend = 3600
                else:
                    addend = -3600
                timeTuple = time.localtime(t + addend)
        # dfn = destination filename
        dfn = self.rotation_filename(self.baseFilename + "." +
                                     time.strftime(self.suffix, timeTuple))
        if os.path.exists(dfn):
            # Already rolled over.
            return

        if self.stream:
            self.stream.close()
            self.stream = None
        self.rotate(self.baseFilename, dfn)
        if self.backupCount > 0:
            for s in self.getFilesToDelete():
                os.remove(s)
        if not self.delay:
            self.stream = self._open()
        self.rolloverAt = self.computeRollover(currentTime)

        # compress the log file
        self.gzip_logs(dfn)
