from __future__ import annotations
from logging import FileHandler
from queue import Queue
from threading import Thread, Event, Lock
import time
from queue import Empty
import os
from pathlib import Path
import zstandard as zstd
import py7zr
from py7zr import FILTER_LZMA2
import shutil
from enum import Enum
import traceback

class LogCompressorType(Enum):
    NONE = 0
    ZSTD = 1
    SEVEN_ZIP = 2

class AsyncFileHandler(FileHandler):
    """Asynchronous file handler that compresses log files."""
    compressor_type = LogCompressorType.NONE
    @classmethod
    def set_compressor_type(cls, compressor_type: LogCompressorType):
        cls.compressor_type = compressor_type

    is_del_log_after_compress = True
    @classmethod
    def set_del_log_after_compress(cls, del_log_after_compress: bool):
        cls.is_del_log_after_compress = del_log_after_compress

    def __init__(self, filename, mode='a', encoding='utf-8', delay=False):
        super().__init__(filename, mode, encoding, delay)
        self.queue = Queue()
        self._stop = Event()
        self.thread = Thread(target=self.consume_log)
        self.thread.daemon = True
        self.thread.start()
        self.emit_lock = Lock()
        self.emit_done_event = Event()
        if not delay:
            folder = Path(os.path.dirname(self.baseFilename))
            folder.mkdir(parents=True, exist_ok=True)
        self.delay = delay
        self.compress_filepath: Path | None = None

    def _open(self):
        if self.delay:
            folder = Path(os.path.dirname(self.baseFilename))
            folder.mkdir(parents=True, exist_ok=True)
            self.compress_filepath = None
        return super()._open()

    def consume_log(self):
        while True:
            if self.queue.empty():
                if self._stop.is_set():
                    break
                time.sleep(0.001)
                continue
            self.emit_done_event.clear()
            try:
                record = self.queue.get(timeout=0.001)
            except Empty:
                continue
            with self.emit_lock:
                super().emit(record)
            self.emit_done_event.set()

    def wait_queue_writing(self):
        if self.queue.empty():
            return
        while True:
            self.emit_done_event.wait(timeout=0.001)
            if self.queue.empty():
                return
            time.sleep(0.001)

    def emit(self, record):
        self.queue.put(record)

    def close(self):
        self._stop.set()
        self.thread.join()
        super().close()
        self.compress_log()
        
    def compress_log(self):
        if self.compressor_type == LogCompressorType.NONE:
            return
        if self.compressor_type == LogCompressorType.ZSTD:
            self.compress_log_zstd()
        elif self.compressor_type == LogCompressorType.SEVEN_ZIP:
            self.compress_log_7z()
        if self.is_del_log_after_compress:
            Path(self.baseFilename).unlink(missing_ok=True)

    def compress_log_zstd(self):
        try:
            src = Path(self.baseFilename)
            if src.exists() and src.stat().st_size > 0:
                dst = src.with_suffix(src.suffix + ".zst")
                cctx = zstd.ZstdCompressor(level=5)
                with src.open("rb") as fi, dst.open("wb") as fo:
                    with cctx.stream_writer(fo) as compressor:
                        shutil.copyfileobj(fi, compressor)
                self.compress_filepath = Path(dst)
        except Exception:
            print("Compress log to zstd failed")

    def compress_log_7z(self):
        try:
            src = Path(self.baseFilename)
            if src.exists() and src.stat().st_size > 0:
                dst = src.with_suffix(src.suffix + ".7z")
                filters = [{'id': FILTER_LZMA2, 'preset': 5}]
                with py7zr.SevenZipFile(dst, mode='w', filters=filters) as z:
                    z.write(src, arcname=src.name)
                self.compress_filepath = Path(dst)
        except Exception:
            print("Compress log to 7z failed")
    
    def switch_folder(self, new_folder):
        new_folder = Path(new_folder)
        new_folder.mkdir(parents=True, exist_ok=True)
        new_filename = os.path.join(new_folder,
                                    os.path.basename(self.baseFilename))
        self.switch_file(new_filename)

    def switch_file(self, new_filename):
        self.wait_queue_writing()
        with self.emit_lock:
            if self.stream:
                self.stream.close()
                self.compress_log()
            self.baseFilename = self._resolve_filename(new_filename)
            self.stream = self._open()
    
    def move_file(self, new_filename):
        with self.emit_lock:
            self.flush()
            new_filename = Path(self._resolve_filename(new_filename))
            new_filename.parent.mkdir(parents=True, exist_ok=True)
            if self.stream:
                self.stream.close()
                os.rename(self.baseFilename, new_filename)
            self.baseFilename = str(new_filename)
            self.stream = self._open()

    def get_log_folder(self) -> Path:
        return Path(self.baseFilename).parent

    def copy_file(self, target_filename):
        with self.emit_lock:
            self.flush()
            target_filename = Path(self._resolve_filename(target_filename))
            if self.stream:
                self.stream.close()
                target_filename.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy(self.baseFilename, target_filename)
                print(f"Copy log to {target_filename}")
            self.stream = self._open()
    
    def _resolve_filename(self, filename):
        return os.path.abspath(filename)