# ruff: noqa: E402
# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

# lance dataset distributed transform job checkpointing + UDF utils

import base64
import fcntl
import json
import logging
import os
import site
import tempfile
import threading
import zipfile
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

_LOG = logging.getLogger(__name__)

_extract_lock = threading.Lock()


def _download_and_extract(args: dict) -> None:
    # Import pyarrow lazily to avoid issues when pyarrow isn't available yet
    # (e.g., when it's part of the manifest being loaded via GENEVA_ZIPS)
    import pyarrow.fs as fs

    src_path, download_path, output_dir = (
        args["src_path"],
        args["download_path"],
        args["output_dir"],
    )
    handle, path = fs.FileSystem.from_uri(src_path)
    handle: fs.FileSystem = handle
    path: str = path

    with (
        handle.open_input_file(path) as f,
        open(download_path, "wb") as out,
    ):
        chunk_size = 1024 * 1024 * 8  # 8MiB chunks
        while data := f.read(chunk_size):
            out.write(data)

    with (
        zipfile.ZipFile(download_path) as z,
        _extract_lock,  # ensure only one thread extracts at a time
    ):
        z.extractall(output_dir)
        _LOG.info("extracted workspace to %s", output_dir)


# MAGIC: if GENEVA_ZIPS is set, we will extract the zips and add them as site-packages
# this is how we acheive "import geneva" == importing workspace from client
#
# NOTE: think of this like booting up a computer. At this point we do not have any
# dependencies installed, so this logic needs to have minimal dependency surface.
# We avoid importing anything from geneva and do everything in the stdlib
if "GENEVA_ZIPS" in os.environ:
    import fcntl

    with (
        open("/tmp/.geneva_zip_setup", "w") as file,  # noqa: S108
        ThreadPoolExecutor(max_workers=8) as executor,
    ):
        # use fcntl to lock the file so we don't have multiple processes
        # trying to extract at the same time and blow up the disk space
        fcntl.lockf(file, fcntl.LOCK_EX)

        payload = json.loads(base64.b64decode(os.environ["GENEVA_ZIPS"]))
        zips = payload.get("zips", [])

        for parts in zips:
            if not len(parts):
                # got an empty list, skip
                continue

            _LOG.info("Setting up geneva workspace from zips %s", parts)
            file_name = parts[0].split("/")[-1]
            name = file_name.split(".")[0]
            output_dir = Path(tempfile.gettempdir()) / name
            if output_dir.exists():
                _LOG.info("workspace already extracted to %s", output_dir)
            else:
                # force collect to surface errors
                list(
                    executor.map(
                        _download_and_extract,
                        (
                            {
                                "src_path": z,
                                "download_path": Path(tempfile.gettempdir())
                                / z.split("/")[-1],
                                "output_dir": output_dir,
                            }
                            for z in parts
                        ),
                    )
                )

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

        fcntl.lockf(file, fcntl.LOCK_UN)


from geneva._context import get_current_context
from geneva.apply import CheckpointingApplier, ReadTask, ScanTask
from geneva.checkpoint import (
    CheckpointStore,
    InMemoryCheckpointStore,
)
from geneva.db import connect
from geneva.transformer import udf

__all__ = [
    "CheckpointStore",
    "connect",
    "InMemoryCheckpointStore",
    "CheckpointingApplier",
    "ReadTask",
    "ScanTask",
    "udf",
    "get_current_context",
]

version = "0.7.0"

__version__ = version
