from distributed.diagnostics.plugin import SchedulerPlugin,WorkerPlugin
from dask.utils import tmpfile
import os
import sys
import uuid
import zipfile

#Modeled after UploadDirectory nanny plugin.
#Needed since ~2023 Dask serialize/deserialize architecture now requires scheduler and workers to have dependencies present.
class UploadDirectory_Base():
    """A plugin template to upload a local directory to scheduler/workers.

    Parameters
    ----------
    path: str
        A path to the directory to upload
    skip_words:
        Sub folders to ignore
    skip:
        files to ignores
    update_path:
        Update Python path with uploaded directory.
    """

    def __init__(
        self, 
        path,
        skip_words=(".git", ".github", ".pytest_cache", "tests", "docs"),
        skip=(lambda fn: os.path.splitext(fn)[1] == ".pyc",),
        update_path=False,
    ):
        """
        Initialize the plugin mixin by reading in the data from the given file.
        """
        path = os.path.expanduser(path)
        self.path = os.path.split(path)[-1]
        self.name = "upload-directory-" + os.path.split(path)[-1]
        self.update_path = update_path

        with tmpfile(extension="zip") as fn:
            with zipfile.ZipFile(fn, "w", zipfile.ZIP_DEFLATED) as z:
                for root, dirs, files in os.walk(path):
                    for file in files:
                        filename = os.path.join(root, file)
                        if any(predicate(filename) for predicate in skip):
                            continue
                        dirs = filename.split(os.sep)
                        if any(word in dirs for word in skip_words):
                            continue

                        archive_name = os.path.relpath(
                            os.path.join(root, file), os.path.join(path, "..")
                        )
                        z.write(filename, archive_name)

            with open(fn, "rb") as f:
                self.data = f.read()
    
    async def extract(self, local_directory):
        """
        Extracts the bufferized and zipped folder into a local directory.
        """
        fn = os.path.join(local_directory, f"tmp-{uuid.uuid4()}.zip")
        with open(fn, "wb") as f:
            f.write(self.data)

        import zipfile

        tmp_dir = os.path.join(local_directory, f"tmp-{uuid.uuid4()}")
        os.mkdir(tmp_dir)
        with zipfile.ZipFile(fn) as z:
            z.extractall(path=tmp_dir)

        if self.update_path:
            path = tmp_dir
            #Add local_directory/tmp_dir to path
            if path not in sys.path:
                sys.path.insert(0, path)
            #Add local_directory/tmp_dir/{upload directory name} to path
            if os.path.join(tmp_dir, self.path) not in sys.path:
                sys.path.insert(0, os.path.join(tmp_dir, self.path))

        os.remove(fn)

class UploadDirectory_Scheduler(UploadDirectory_Base, SchedulerPlugin):
    async def start(self, scheduler):
        await self.extract(scheduler.local_directory)

class UploadDirectory_Worker(UploadDirectory_Base, WorkerPlugin):
    async def setup(self, worker):
        await self.extract(worker.local_directory)

import cloudpickle
from . import dask_plugins
cloudpickle.register_pickle_by_value(dask_plugins)

...

plugin = dask_plugins.UploadDirectory_Scheduler(os.getcwd(), update_path=True)
client.register_plugin(plugin)
plugin = dask_plugins.UploadDirectory_Worker(os.getcwd(), update_path=True)
client.register_plugin(plugin)