# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Geneva Authors

# a zip packager for local workspace

import contextlib
import hashlib
import logging
import queue
import re
import site
import sys
import tempfile
import zipfile
from pathlib import Path

import attrs
import emoji
import multiprocess
from multiprocess.queues import Queue

from geneva.config import ConfigBase
from geneva.tqdm import tqdm

_LOG = logging.getLogger(__name__)


@attrs.define
class _ZipperConfig(ConfigBase):
    output_path: Path | None = attrs.field(
        validator=attrs.validators.instance_of(Path),
        converter=attrs.converters.optional(Path),
    )

    @classmethod
    def name(cls) -> str:
        return "zipper"


_DEFAULT_IGNORES = [
    r"\.pyc",
    r".*__pycache__.*",
    r"\.venv",
    r"\.git",
    r"\.ruff_cache",
    r"\.vscode",
    r"\.github",
    r"\.lance",
]


@attrs.define
class _ProgressQueue:
    queue: Queue = attrs.field()

    def update(self, value: int) -> None:
        self.queue.put(value)

    def get(self, *args, **kwargs) -> int:
        return self.queue.get(*args, **kwargs)


@attrs.define
class WorkspaceZipper:
    path: Path = attrs.field(
        converter=attrs.converters.pipe(
            Path,
            Path.resolve,
            Path.absolute,
        )
    )

    @path.validator
    def _path_validator(self, attribute, value: Path) -> None:
        if not value.is_dir():
            raise ValueError("path must be a directory")

        # make sure the path is the current working directory, or
        # is part of sys.path
        if value == Path.cwd().resolve().absolute():
            return

        sys_paths = {Path(x).resolve().absolute() for x in sys.path}

        if value not in sys_paths:
            raise ValueError("path must be cwd or part of sys.path")

    output_dir: Path = attrs.field(converter=Path)

    @output_dir.default
    def _output_dir_default(self) -> Path:
        config = _ZipperConfig.get()
        if config.output_path is not None:
            return config.output_path
        return self.path / ".geneva"

    ignore_regexs: list[re.Pattern] = attrs.field(
        factory=lambda: [re.compile(r) for r in _DEFAULT_IGNORES],
        converter=lambda x: [re.compile(r) for r in x],
    )
    """
    a list of regex patterns to ignore when zipping the workspace

    only ignores based on the relative path of the file
    """

    shard_size: int = attrs.field(
        default=128 * 1024**2,
    )
    """
    the size of each shard in bytes
    """

    def checksum(self) -> tuple[str, int, list[tuple[str, int]]]:
        """
        create a checksum for the files in the workspace based on name and mtime.

        return the path of the zip file and a content hash of the zip file
        """

        with tqdm(self.path.rglob("*")) as pbar:
            total_size = 0
            pbar.set_description(
                emoji.emojize(
                    f":magnifying_glass_tilted_left: scanning workspace: {self.path}"
                )
            )

            hasher = hashlib.sha256()

            files = []
            for child in pbar:
                child: Path = child
                # Path.relative_to() is too slow
                # here is the benchmark:
                # removeprefix is 10x faster than relative_to
                #
                # In [49]: %timeit for p in cwd.rglob("*"): p.as_posix().removeprefix(cwd.as_posix() + "/") # noqa: E501
                # 134 ms ± 663 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
                #
                # In [50]: %timeit for p in cwd.rglob("*"): p.relative_to(cwd)
                # 1.32 s ± 22.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
                arcname = child.as_posix().removeprefix(self.path.as_posix() + "/")
                if any(r.match(arcname) for r in self.ignore_regexs):
                    continue
                stat = child.stat()
                total_size += stat.st_size

                files.append((arcname, stat.st_size))
                hasher.update(arcname.encode("utf-8"))
                hasher.update(stat.st_mtime_ns.to_bytes(length=8, byteorder="little"))
            checksum = hasher.hexdigest()
            pbar.set_description(
                emoji.emojize(
                    f":white_check_mark: scanned workspace: checksum={checksum}"
                )
            )
            return checksum, total_size, files

    def zip(
        self, create_zips=True, excludes: list[Path] | None = None
    ) -> tuple[list[Path], str]:
        """
        create a zip file for the workspace

        return the path of the zip file and a content hash of the zip file
        """
        if excludes is None:
            excludes = []
        checksum, total_size, files = self.checksum()

        with tqdm(
            total=total_size, unit="B", unit_scale=True, unit_divisor=1024, smoothing=0
        ) as pbar:
            if create_zips:
                pbar.set_description(
                    emoji.emojize(f":card_file_box: zipping workspace: {self.path}")
                )
            else:
                pbar.set_description(
                    emoji.emojize(
                        f":card_file_box: zipping workspace: {self.path} (dry run)"
                    )
                )

            # just a single zip if the size is small
            if total_size <= self.shard_size:
                zip_path = self.output_dir / f"{checksum}.zip"
                _LOG.info(f"Zipping workspace {self.path} to {zip_path}")
                if zip_path in excludes:
                    _LOG.info("Zip exists, skipping: %s", zip_path)
                    return [], ""

                if create_zips:
                    _LOG.info(f"creating zip {zip_path} with {len(files)} files")
                    self._write_zip(zip_path, files, pbar)
                    zip_path.chmod(0o777)
                    zip_path.parent.chmod(0o777)
                return [zip_path], checksum

            # workspace is large, shard it into multiple zip files
            def _chunk_by_size():  # noqa: ANN202
                running_size = 0
                current_chunk = []
                for arcname, size in files:
                    current_chunk.append((arcname, size))
                    running_size += size
                    if running_size >= self.shard_size:
                        yield current_chunk
                        current_chunk = []
                        running_size = 0
                if current_chunk:
                    yield current_chunk

            ctx = multiprocess.context.SpawnContext()
            with (
                ctx.Pool() as pool,
                ctx.Manager() as manager,
            ):
                progress_queue = _ProgressQueue(manager.Queue())

                chunks = [
                    (
                        self.output_dir / f"{checksum}.part{idx:02d}.zip",
                        file_and_sizes,
                        progress_queue,
                    )
                    for idx, file_and_sizes in enumerate(_chunk_by_size())
                ]

                if create_zips:
                    _LOG.info(f"{len(chunks)} chunks to zip")
                    filtered_chunks = [
                        (zip_name, fands, pq)
                        for zip_name, fands, pq in chunks
                        if zip_name not in excludes
                    ]
                    _LOG.info(f"{len(filtered_chunks)} filtered chunks to zip:")
                    if len(filtered_chunks) == 0:
                        # all filtered out, return empty list.
                        return [], ""

                    join_handle = pool.starmap_async(
                        self._write_zip,
                        filtered_chunks,
                    )

                    while pbar.n < total_size:
                        # join handle is ready before we expect it to
                        # this means there was an error in the worker
                        # get the future to raise the error
                        if join_handle.ready():
                            join_handle.get()

                        with contextlib.suppress(queue.Empty):
                            pbar.update(progress_queue.get(block=False))

                    join_handle.get()

            return [chunk[0] for chunk in chunks], checksum

    def _write_zip(
        self,
        zip_path: Path,
        file_and_sizes: list[tuple[str, int]],
        pbar,
    ) -> None:
        if zip_path.exists():
            _LOG.info("Zip exists, skipping: %s", zip_path)
            return

        _LOG.info(f"Zipping workspace {self.path} to {zip_path}")

        # use compression level 1 to speed up the process
        # the output size is not too different from level 6
        with zipfile.ZipFile(
            zip_path, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=1
        ) as z:
            for arcname, size in file_and_sizes:
                z.write(self.path / arcname, arcname)
                pbar.update(size)


@attrs.define
class _UnzipperConfig(ConfigBase):
    output_dir: Path | None = attrs.field(
        converter=attrs.converters.optional(Path),
        default=None,
    )

    @classmethod
    def name(cls) -> str:
        return "unzipper"


# Kept separate from the zipper to avoid config mess
@attrs.define
class WorkspaceUnzipper:
    output_dir: Path = attrs.field(
        converter=attrs.converters.pipe(
            attrs.converters.default_if_none(
                _UnzipperConfig.get().output_dir,
            ),
            attrs.converters.default_if_none(
                factory=tempfile.mkdtemp,
            ),
            Path,
            Path.resolve,
            Path.absolute,
        ),
        default=None,
    )

    def unzip(self, zip_path: Path) -> None:
        """
        extract the zip file to the workspace
        """
        with zipfile.ZipFile(zip_path, "r") as z:
            z.extractall(self.output_dir)
        _LOG.info("extracted workspace to %s", self.output_dir)

        site.addsitedir(self.output_dir.as_posix())
        _LOG.info("added %s to sys.path", self.output_dir)
