import os
import shutil
import threading
import queue
import time
import logging
from uuid import uuid4 as uuid
from collections import namedtuple, defaultdict
from pymrf4.utils import set_interval

logger = logging.getLogger(__name__)

FileCopyTask = namedtuple("FileCopyTask", ["tag", "src", "dest"])

class Harvestor:
    # implement a singleton pattern
    _instance = None


    def __init__(self, num_of_workers) -> None:
        logger.info("Harvestor initialized with num_of_workers: %d", num_of_workers)
        self.num_of_workers: int = num_of_workers
        self.workers: list[FileCopyWorker] = []
        # self.event: threading.Event = threading.Event()
        self.queue = FileCopyQueue()


    def setup(harvestor_config):
        logger.info(f"Harvestor setup with config: {harvestor_config}")
        Harvestor._instance = Harvestor(harvestor_config.get("workers", 3))


    def start(self, src_dirs, dest_dir, tag=None):
        if tag is None:
            tag = uuid()

        logger.info(f"Harvestor start with src_dirs: {src_dirs}, dest_dir: {dest_dir} tag: {tag}")

        if self.queue.unfinished_tasks == 0:
            self.queue.reset()

        self.queue.add_files(src_dirs, dest_dir, tag)
        # self.event.set()

        dead_workers: list[FileCopyWorker] = []
        for worker in self.workers:
            if not worker.is_alive():
                dead_workers.append(worker)

        for worker in dead_workers:
            self.workers.remove(worker)

        if len(self.workers) < self.num_of_workers:
            new_workers = [FileCopyWorker(self.queue) for _ in range(self.num_of_workers - len(self.workers))]
            self.workers.extend(new_workers)
            for thread in new_workers:
                thread.start()

        # set interval for heartbeat
        self.heartbeat_timer = set_interval(self.heartbeat, 11)


    @staticmethod
    def harvest(src_dir, dest_dir, tag=None):
        # check if there's active workers
        if Harvestor._instance is None:
            Harvestor._instance = Harvestor()

        Harvestor._instance.start([src_dir], dest_dir, tag)

    @staticmethod
    def stop():
        if Harvestor._instance is not None:
            Harvestor._instance.heartbeat_timer.cancel()

            for worker in Harvestor._instance.workers:
                worker.stop()

            Harvestor._instance.queue.reset()


    @staticmethod
    def is_working():
        return Harvestor._instance is not None and \
            Harvestor._instance.queue.unfinished_tasks > 0


    def heartbeat(self):
        pass
        # for worker in self.workers:
        #     if not worker.is_alive():
        #         logger.error(f"Worker {worker.name} is dead.")
        #         self.workers.remove(worker)
        #         # restart the worker
        #         new_worker = FileCopyWorker(self.queue)
        #         self.workers.append(new_worker)
        #         new_worker.start()


class FileCopyQueue(queue.Queue):
    REPORT_SPEED_INTERVAL: int = 1 # second

    def __init__(self):
        super().__init__()
        self.reset()

    def reset(self):
        self.all_src_dirs = set()
        self.tags = defaultdict(int)
        self.seen_files = set()
        self.start_time = time.time()
        self.total_size = 0
        self.copied_size = 0
        self.transfer_speed = 0
        self.last_report_time = self.start_time
        self.copied_size_since_last_report = 0

        # Queue's
        self.unfinished_tasks = 0


    def add_files(self, src_dirs, dest_dir, tag):
        files_added = 0
        batch_size = 0
        for src_dir in src_dirs:
            if os.path.isfile(src_dir):
                # single file torrents
                wrap_dir = os.path.basename(os.path.splitext(src_dir)[0]).strip()
                dest_file = os.path.join(dest_dir, wrap_dir, os.path.basename(src_dir))

                task = FileCopyTask(tag, src_dir, dest_file)
                src_size = self.__add_file_copy_task(task)

                files_added += 1
                batch_size += src_size
                # self.put((src_dir, dest_file))  # Queue's
            else:
                src_dir_name = os.path.basename(src_dir)
                for root, _, files in os.walk(src_dir):
                    for file in files:
                        src_file = os.path.join(root, file)
                        dest_file = os.path.join(dest_dir, src_dir_name, os.path.relpath(src_file, src_dir))

                        task = FileCopyTask(tag, src_file, dest_file)
                        src_size = self.__add_file_copy_task(task)

                        files_added += 1
                        batch_size += src_size

                        # if (src_file, dest_file) not in self.seen_files:
                        #     self.seen_files.add((src_file, dest_file))
                        #     # self.put((src_file, dest_file)) # Queue's
                        #     self.put(FileCopyTask(tag, src_file, dest_file)) # Queue's
                        #     self.total_size += os.path.getsize(src_file)
                        #     self.tags[tag] += 1

                self.all_src_dirs.add(src_dir)

        logger.info(f"{files_added} files added, size: {batch_size}. Total size: {self.total_size}")


    def __add_file_copy_task(self, task: FileCopyTask):
        """ Internal method that add FileCopyTask to queue and update stats
        """
        if (task.src, task.dest) not in self.seen_files:
            self.seen_files.add((task.src, task.dest))
            self.put(task)  # for Queue
            self.total_size += os.path.getsize(task.src)
            self.tags[task.tag] += 1

            return os.path.getsize(task.src)


    # override
    def task_done(self, task) -> None:
        self.tags[task.tag] -= 1
        if (self.tags[task.tag] == 0):
            del self.tags[task.tag]
            # notify harvestor that the task is done

        return super().task_done()

        # to_remove: list[str] = []
        # # check if src_dirs can delete
        # for src_dir in self.all_src_dirs:
        #     if os.path.isdir(src_dir) and len(os.listdir(src_dir)) == 0:
        #         shutil.rmtree(src_dir)
        #         to_remove.append(src_dir)
        # for dir in to_remove:
        #     self.all_src_dirs.remove(dir)

    def update_progress(self, bytes_copied):
        self.copied_size += bytes_copied
        self.copied_size_since_last_report += bytes_copied

        # update speed
        current_time = time.time()
        if current_time - self.last_report_time >= FileCopyQueue.REPORT_SPEED_INTERVAL:
            self.transfer_speed = self.copied_size_since_last_report / (current_time - self.last_report_time)
            self.last_report_time = current_time
            self.copied_size_since_last_report = 0


    def get_progress(self):
        return (self.copied_size / self.total_size) * 100 if self.total_size > 0 else 0


    def get_speed(self):
        # elapsed_time = time.time() - self.start_time
        # return (self.copied_size / elapsed_time) if elapsed_time > 0 else 0
        return self.transfer_speed


class FileCopyWorker(threading.Thread):
    """Consumes the file queue and copies files from src to dest.
    """
    def __init__(self, task_queue: FileCopyQueue):
        super().__init__()
        self.name = "FileCopyWorker"
        self.task_queue = task_queue
        # self.event = event
        self.should_stop = False

    # def run(self):
    #     while True:
    #         self.event.wait(30)
    #         if not self.event.is_set():
    #             continue

    #         try:
    #             while True:
    #                 if self.should_stop:
    #                     self.event.clear()
    #                     break

    #                 src, dest = self.task_queue.get(block=False)
    #                 self.copy_file(src, dest)
    #                 # delete src file
    #                 os.remove(src)

    #                 self.task_queue.task_done()
    #         except queue.Empty:
    #             self.event.clear()
    #         finally:
    #             self.event.clear()

    def run(self):
        while True:
            if self.should_stop:
                break
            try:
                task = self.task_queue.get(block=False)
                _, src, dest = task
                logger.info(f"Copying {src} to {dest}")
                if self.__copy_file(src, dest):
                    # delete src file only when copy is successful
                    os.remove(src)
                    logger.info(f"Copied {src} to {dest}")
                    self.task_queue.task_done(task)

            except queue.Empty:
                break
            else:
                logger.debug(f"Queue has {self.task_queue.unfinished_tasks} tasks left.")


    def __copy_file(self, src, dest):
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        buffer_size = 32 * 1024 * 1024  # buffer size

        with open(src, 'rb') as fsrc, open(dest, 'wb') as fdest:
            while chunk := fsrc.read(buffer_size):
                if self.should_stop:
                    return False
                fdest.write(chunk)
                self.task_queue.update_progress(len(chunk))

        return True

    def stop(self):
        self.should_stop = True


# class ProgressReporter(threading.Thread):
#     def __init__(self, task_queue):
#         super().__init__()
#         self.task_queue = task_queue

#     def run(self):
#         start_time = time.time()
#         while self.task_queue.copied_size < self.task_queue.total_size:
#             time.sleep(1)
#             progress = self.task_queue.get_progress()
#             speed = self.task_queue.get_speed(start_time)
#             print(f"Progress: {progress:.2f}%, Speed: {speed:.2f} MB/s")


# if __name__ == "__main__":
#     task_queue = FileCopyQueue()
#     src_dirs = ['/Users/xhguo/Downloads/!故事机', '/Users/xhguo/Downloads/01. Axure 9 零基础到实战(80节视频+课程元件)']
#     target_dir = '/Volumes/GN13/Target'

#     print("File transfer started.")

#     task_queue.add_files(src_dirs, target_dir)

#     # Start file copy worker threads
#     workers = [FileCopyWorker(task_queue) for _ in range(6)]
#     for thread in workers:
#         thread.start()

#     progress_reporter = ProgressReporter(task_queue)
#     progress_reporter.start()

#     for worker in workers:
#         worker.join()
#     progress_reporter.join()

#     print("File transfer completed.")