import atexit
import codecs
import errno
import fnmatch
import gzip
import json
import logging
import math
import os
import pathlib
import posixpath
import shutil
import stat
import subprocess
import sys
import tarfile
import tempfile
import time
import urllib.parse
import urllib.request
import uuid
from concurrent.futures import as_completed
from contextlib import contextmanager
from dataclasses import dataclass
from subprocess import CalledProcessError, TimeoutExpired
from typing import Optional, Union
from urllib.parse import unquote
from urllib.request import pathname2url

import yaml

from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE

try:
    from yaml import CSafeDumper as YamlSafeDumper
    from yaml import CSafeLoader as YamlSafeLoader
except ImportError:
    from yaml import SafeDumper as YamlSafeDumper
    from yaml import SafeLoader as YamlSafeLoader

from mlflow.entities import FileInfo
from mlflow.environment_variables import (
    _MLFLOW_MPD_NUM_RETRIES,
    _MLFLOW_MPD_RETRY_INTERVAL_SECONDS,
    MLFLOW_DOWNLOAD_CHUNK_TIMEOUT,
    MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR,
)
from mlflow.exceptions import MissingConfigException, MlflowException
from mlflow.protos.databricks_artifacts_pb2 import ArtifactCredentialType
from mlflow.utils import download_cloud_file_chunk, merge_dicts
from mlflow.utils.databricks_utils import _get_dbutils
from mlflow.utils.os import is_windows
from mlflow.utils.process import cache_return_value_per_process
from mlflow.utils.request_utils import cloud_storage_http_request, download_chunk
from mlflow.utils.rest_utils import augmented_raise_for_status

ENCODING = "utf-8"
MAX_PARALLEL_DOWNLOAD_WORKERS = os.cpu_count() * 2
_PROGRESS_BAR_DISPLAY_THRESHOLD = 500_000_000  # 500 MB

_logger = logging.getLogger(__name__)


_logger = logging.getLogger(__name__)


class ArtifactProgressBar:
    def __init__(self, desc, total, step, **kwargs) -> None:
        self.desc = desc
        self.total = total
        self.step = step
        self.pbar = None
        self.progress = 0
        self.kwargs = kwargs

    def set_pbar(self):
        if MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR.get():
            try:
                from tqdm.auto import tqdm

                self.pbar = tqdm(total=self.total, desc=self.desc, **self.kwargs)
            except ImportError:
                pass

    @classmethod
    def chunks(cls, file_size, desc, chunk_size):
        bar = cls(
            desc,
            total=file_size,
            step=chunk_size,
            unit="iB",
            unit_scale=True,
            unit_divisor=1024,
            miniters=1,
        )
        if file_size >= _PROGRESS_BAR_DISPLAY_THRESHOLD:
            bar.set_pbar()
        return bar

    @classmethod
    def files(cls, desc, total):
        bar = cls(desc, total=total, step=1)
        bar.set_pbar()
        return bar

    def update(self):
        if self.pbar:
            update_step = min(self.total - self.progress, self.step)
            self.pbar.update(update_step)
            self.pbar.refresh()
            self.progress += update_step

    def __enter__(self):
        return self

    def __exit__(self, *args):
        if self.pbar:
            self.pbar.close()


def is_directory(name):
    return os.path.isdir(name)


def is_file(name):
    return os.path.isfile(name)


def exists(name):
    return os.path.exists(name)


def list_all(root, filter_func=lambda x: True, full_path=False):
    """
    List all entities directly under 'dir_name' that satisfy 'filter_func'

    :param root: Name of directory to start search
    :param filter_func: function or lambda that takes path
    :param full_path: If True will return results as full path including `root`

    :return: list of all files or directories that satisfy the criteria.
    """
    if not is_directory(root):
        raise Exception(f"Invalid parent directory '{root}'")
    matches = [x for x in os.listdir(root) if filter_func(os.path.join(root, x))]
    return [os.path.join(root, m) for m in matches] if full_path else matches


def list_subdirs(dir_name, full_path=False):
    """
    Equivalent to UNIX command:
      ``find $dir_name -depth 1 -type d``

    :param dir_name: Name of directory to start search
    :param full_path: If True will return results as full path including `root`

    :return: list of all directories directly under 'dir_name'
    """
    return list_all(dir_name, os.path.isdir, full_path)


def list_files(dir_name, full_path=False):
    """
    Equivalent to UNIX command:
      ``find $dir_name -depth 1 -type f``

    :param dir_name: Name of directory to start search
    :param full_path: If True will return results as full path including `root`

    :return: list of all files directly under 'dir_name'
    """
    return list_all(dir_name, os.path.isfile, full_path)


def find(root, name, full_path=False):
    """
    Search for a file in a root directory. Equivalent to:
      ``find $root -name "$name" -depth 1``

    :param root: Name of root directory for find
    :param name: Name of file or directory to find directly under root directory
    :param full_path: If True will return results as full path including `root`

    :return: list of matching files or directories
    """
    path_name = os.path.join(root, name)
    return list_all(root, lambda x: x == path_name, full_path)


def mkdir(root, name=None):
    """
    Make directory with name "root/name", or just "root" if name is None.

    :param root: Name of parent directory
    :param name: Optional name of leaf directory

    :return: Path to created directory
    """
    target = os.path.join(root, name) if name is not None else root
    try:
        os.makedirs(target)
    except OSError as e:
        if e.errno != errno.EEXIST or not os.path.isdir(target):
            raise e
    return target


def make_containing_dirs(path):
    """
    Create the base directory for a given file path if it does not exist; also creates parent
    directories.
    """
    dir_name = os.path.dirname(path)
    if not os.path.exists(dir_name):
        os.makedirs(dir_name)


def write_yaml(root, file_name, data, overwrite=False, sort_keys=True, ensure_yaml_extension=True):
    """
    Write dictionary data in yaml format.

    :param root: Directory name.
    :param file_name: Desired file name.
    :param data: data to be dumped as yaml format
    :param overwrite: If True, will overwrite existing files
    :param ensure_yaml_extension: If True, will automatically add .yaml extension if not given
    """
    if not exists(root):
        raise MissingConfigException(f"Parent directory '{root}' does not exist.")

    file_path = os.path.join(root, file_name)
    yaml_file_name = file_path
    if ensure_yaml_extension and not file_path.endswith(".yaml"):
        yaml_file_name = file_path + ".yaml"

    if exists(yaml_file_name) and not overwrite:
        raise Exception(f"Yaml file '{file_path}' exists as '{yaml_file_name}")

    try:
        with codecs.open(yaml_file_name, mode="w", encoding=ENCODING) as yaml_file:
            yaml.dump(
                data,
                yaml_file,
                default_flow_style=False,
                allow_unicode=True,
                sort_keys=sort_keys,
                Dumper=YamlSafeDumper,
            )
    except Exception as e:
        raise e


def overwrite_yaml(root, file_name, data, ensure_yaml_extension=True):
    """
    Safely overwrites a preexisting yaml file, ensuring that file contents are not deleted or
    corrupted if the write fails. This is achieved by writing contents to a temporary file
    and moving the temporary file to replace the preexisting file, rather than opening the
    preexisting file for a direct write.

    :param root: Directory name.
    :param file_name: File name.
    :param data: The data to write, represented as a dictionary.
    :param ensure_yaml_extension: If True, Will automatically add .yaml extension if not given
    """
    tmp_file_path = None
    original_file_path = os.path.join(root, file_name)
    original_file_mode = os.stat(original_file_path).st_mode
    try:
        tmp_file_fd, tmp_file_path = tempfile.mkstemp(suffix="file.yaml")
        os.close(tmp_file_fd)
        write_yaml(
            root=get_parent_dir(tmp_file_path),
            file_name=os.path.basename(tmp_file_path),
            data=data,
            overwrite=True,
            sort_keys=True,
            ensure_yaml_extension=ensure_yaml_extension,
        )
        shutil.move(tmp_file_path, original_file_path)
        # restores original file permissions, see https://docs.python.org/3/library/tempfile.html#tempfile.mkstemp
        os.chmod(original_file_path, original_file_mode)
    finally:
        if tmp_file_path is not None and os.path.exists(tmp_file_path):
            os.remove(tmp_file_path)


def read_yaml(root, file_name):
    """
    Read data from yaml file and return as dictionary

    :param root: Directory name
    :param file_name: File name. Expects to have '.yaml' extension

    :return: Data in yaml file as dictionary
    """
    if not exists(root):
        raise MissingConfigException(
            f"Cannot read '{file_name}'. Parent dir '{root}' does not exist."
        )

    file_path = os.path.join(root, file_name)
    if not exists(file_path):
        raise MissingConfigException(f"Yaml file '{file_path}' does not exist.")
    try:
        with codecs.open(file_path, mode="r", encoding=ENCODING) as yaml_file:
            return yaml.load(yaml_file, Loader=YamlSafeLoader)
    except Exception as e:
        raise e


class UniqueKeyLoader(YamlSafeLoader):
    def construct_mapping(self, node, deep=False):
        mapping = set()
        for key_node, _ in node.value:
            key = self.construct_object(key_node, deep=deep)
            if key in mapping:
                raise ValueError(f"Duplicate '{key}' key found in YAML.")
            mapping.add(key)
        return super().construct_mapping(node, deep)


def render_and_merge_yaml(root, template_name, context_name):
    """
    Renders a Jinja2-templated YAML file based on a YAML context file, merge them, and return
    result as a dictionary.

    :param root: Root directory of the YAML files
    :param template_name: Name of the template file
    :param context_name: Name of the context file
    :return: Data in yaml file as dictionary
    """
    from jinja2 import FileSystemLoader, StrictUndefined
    from jinja2.sandbox import SandboxedEnvironment

    template_path = os.path.join(root, template_name)
    context_path = os.path.join(root, context_name)

    for path in (template_path, context_path):
        if not pathlib.Path(path).is_file():
            raise MissingConfigException(f"Yaml file '{path}' does not exist.")

    j2_env = SandboxedEnvironment(
        loader=FileSystemLoader(root, encoding=ENCODING),
        undefined=StrictUndefined,
        line_comment_prefix="#",
    )

    def from_json(input_var):
        with open(input_var, encoding="utf-8") as f:
            return json.load(f)

    j2_env.filters["from_json"] = from_json
    # Compute final source of context file (e.g. my-profile.yml), applying Jinja filters
    # like from_json as needed to load context information from files, then load into a dict
    context_source = j2_env.get_template(context_name).render({})
    context_dict = yaml.load(context_source, Loader=UniqueKeyLoader) or {}

    # Substitute parameters from context dict into template
    source = j2_env.get_template(template_name).render(context_dict)
    rendered_template_dict = yaml.load(source, Loader=UniqueKeyLoader)
    return merge_dicts(rendered_template_dict, context_dict)


def read_parquet_as_pandas_df(data_parquet_path: str):
    """
    Deserialize and load the specified parquet file as a Pandas DataFrame.

    :param data_parquet_path: String, path object (implementing os.PathLike[str]),
    or file-like object implementing a binary read() function. The string
    could be a URL. Valid URL schemes include http, ftp, s3, gs, and file.
    For file URLs, a host is expected. A local file could
    be: file://localhost/path/to/table.parquet. A file URL can also be a path to a
    directory that contains multiple partitioned parquet files. Pyarrow
    support paths to directories as well as file URLs. A directory
    path could be: file://localhost/path/to/tables or s3://bucket/partition_dir.
    :return: pandas dataframe
    """
    import pandas as pd

    return pd.read_parquet(data_parquet_path, engine="pyarrow")


def write_pandas_df_as_parquet(df, data_parquet_path: str):
    """
    Write a DataFrame to the binary parquet format.

    :param df: pandas data frame.
    :param data_parquet_path: String, path object (implementing os.PathLike[str]),
    or file-like object implementing a binary write() function.
    """
    df.to_parquet(data_parquet_path, engine="pyarrow")


class TempDir:
    def __init__(self, chdr=False, remove_on_exit=True):
        self._dir = None
        self._path = None
        self._chdr = chdr
        self._remove = remove_on_exit

    def __enter__(self):
        self._path = os.path.abspath(create_tmp_dir())
        assert os.path.exists(self._path)
        if self._chdr:
            self._dir = os.path.abspath(os.getcwd())
            os.chdir(self._path)
        return self

    def __exit__(self, tp, val, traceback):
        if self._chdr and self._dir:
            os.chdir(self._dir)
            self._dir = None
        if self._remove and os.path.exists(self._path):
            shutil.rmtree(self._path)

        assert not self._remove or not os.path.exists(self._path)
        assert os.path.exists(os.getcwd())

    def path(self, *path):
        return os.path.join("./", *path) if self._chdr else os.path.join(self._path, *path)


def read_file_lines(parent_path, file_name):
    """
    Return the contents of the file as an array where each element is a separate line.

    :param parent_path: Full path to the directory that contains the file.
    :param file_name: Leaf file name.

    :return: All lines in the file as an array.
    """
    file_path = os.path.join(parent_path, file_name)
    with codecs.open(file_path, mode="r", encoding=ENCODING) as f:
        return f.readlines()


def read_file(parent_path, file_name):
    """
    Return the contents of the file.

    :param parent_path: Full path to the directory that contains the file.
    :param file_name: Leaf file name.

    :return: The contents of the file.
    """
    file_path = os.path.join(parent_path, file_name)
    with codecs.open(file_path, mode="r", encoding=ENCODING) as f:
        return f.read()


def get_file_info(path, rel_path):
    """
    Returns file meta data : location, size, ... etc

    :param path: Path to artifact

    :return: `FileInfo` object
    """
    if is_directory(path):
        return FileInfo(rel_path, True, None)
    else:
        return FileInfo(rel_path, False, os.path.getsize(path))


def get_relative_path(root_path, target_path):
    """
    Remove root path common prefix and return part of `path` relative to `root_path`.

    :param root_path: Root path
    :param target_path: Desired path for common prefix removal

    :return: Path relative to root_path
    """
    if len(root_path) > len(target_path):
        raise Exception(f"Root path '{root_path}' longer than target path '{target_path}'")
    common_prefix = os.path.commonprefix([root_path, target_path])
    return os.path.relpath(target_path, common_prefix)


def mv(target, new_parent):
    shutil.move(target, new_parent)


def write_to(filename, data):
    with codecs.open(filename, mode="w", encoding=ENCODING) as handle:
        handle.write(data)


def append_to(filename, data):
    with open(filename, "a") as handle:
        handle.write(data)


def make_tarfile(output_filename, source_dir, archive_name, custom_filter=None):
    # Helper for filtering out modification timestamps
    def _filter_timestamps(tar_info):
        tar_info.mtime = 0
        return tar_info if custom_filter is None else custom_filter(tar_info)

    unzipped_file_handle, unzipped_filename = tempfile.mkstemp()
    try:
        with tarfile.open(unzipped_filename, "w") as tar:
            tar.add(source_dir, arcname=archive_name, filter=_filter_timestamps)
        # When gzipping the tar, don't include the tar's filename or modification time in the
        # zipped archive (see https://docs.python.org/3/library/gzip.html#gzip.GzipFile)
        with gzip.GzipFile(
            filename="", fileobj=open(output_filename, "wb"), mode="wb", mtime=0
        ) as gzipped_tar, open(unzipped_filename, "rb") as tar:
            gzipped_tar.write(tar.read())
    finally:
        os.close(unzipped_file_handle)


def _copy_project(src_path, dst_path=""):
    """
    Internal function used to copy MLflow project during development.

    Copies the content of the whole directory tree except patterns defined in .dockerignore.
    The MLflow is assumed to be accessible as a local directory in this case.


    :param dst_path: MLflow will be copied here
    :return: name of the MLflow project directory
    """

    def _docker_ignore(mlflow_root):
        docker_ignore = os.path.join(mlflow_root, ".dockerignore")
        patterns = []
        if os.path.exists(docker_ignore):
            with open(docker_ignore) as f:
                patterns = [x.strip() for x in f.readlines()]

        def ignore(_, names):
            res = set()
            for p in patterns:
                res.update(set(fnmatch.filter(names, p)))
            return list(res)

        return ignore if patterns else None

    mlflow_dir = "mlflow-project"
    # check if we have project root
    assert os.path.isfile(os.path.join(src_path, "setup.py")), "file not found " + str(
        os.path.abspath(os.path.join(src_path, "setup.py"))
    )
    shutil.copytree(src_path, os.path.join(dst_path, mlflow_dir), ignore=_docker_ignore(src_path))
    return mlflow_dir


def _copy_file_or_tree(src, dst, dst_dir=None):
    """
    :return: The path to the copied artifacts, relative to `dst`
    """
    dst_subpath = os.path.basename(os.path.abspath(src))
    if dst_dir is not None:
        dst_subpath = os.path.join(dst_dir, dst_subpath)
    dst_path = os.path.join(dst, dst_subpath)
    if os.path.isfile(src):
        dst_dirpath = os.path.dirname(dst_path)
        if not os.path.exists(dst_dirpath):
            os.makedirs(dst_dirpath)
        shutil.copy(src=src, dst=dst_path)
    else:
        shutil.copytree(src=src, dst=dst_path, ignore=shutil.ignore_patterns("__pycache__"))
    return dst_subpath


def _get_local_project_dir_size(project_path):
    """
    Internal function for reporting the size of a local project directory before copying to
    destination for cli logging reporting to stdout.
    :param project_path: local path of the project directory
    :return: directory file sizes in KB, rounded to single decimal point for legibility
    """

    total_size = 0
    for root, _, files in os.walk(project_path):
        for f in files:
            path = os.path.join(root, f)
            total_size += os.path.getsize(path)
    return round(total_size / 1024.0, 1)


def _get_local_file_size(file):
    """
    Get the size of a local file in KB
    """
    return round(os.path.getsize(file) / 1024.0, 1)


def get_parent_dir(path):
    return os.path.abspath(os.path.join(path, os.pardir))


def relative_path_to_artifact_path(path):
    if os.path == posixpath:
        return path
    if os.path.abspath(path) == path:
        raise Exception("This method only works with relative paths.")
    return unquote(pathname2url(path))


def path_to_local_file_uri(path):
    """
    Convert local filesystem path to local file uri.
    """
    return pathlib.Path(os.path.abspath(path)).as_uri()


def path_to_local_sqlite_uri(path):
    """
    Convert local filesystem path to sqlite uri.
    """
    path = posixpath.abspath(pathname2url(os.path.abspath(path)))
    prefix = "sqlite://" if sys.platform == "win32" else "sqlite:///"
    return prefix + path


def local_file_uri_to_path(uri):
    """
    Convert URI to local filesystem path.
    No-op if the uri does not have the expected scheme.
    """
    path = uri
    if uri.startswith("file:"):
        parsed_path = urllib.parse.urlparse(uri)
        path = parsed_path.path
        # Fix for retaining server name in UNC path.
        if is_windows() and parsed_path.netloc:
            return urllib.request.url2pathname(rf"\\{parsed_path.netloc}{path}")
    return urllib.request.url2pathname(path)


def get_local_path_or_none(path_or_uri):
    """Check if the argument is a local path (no scheme or file:///) and return local path if true,
    None otherwise.
    """
    parsed_uri = urllib.parse.urlparse(path_or_uri)
    if len(parsed_uri.scheme) == 0 or parsed_uri.scheme == "file" and len(parsed_uri.netloc) == 0:
        return local_file_uri_to_path(path_or_uri)
    else:
        return None


def yield_file_in_chunks(file, chunk_size=100000000):
    """
    Generator to chunk-ify the inputted file based on the chunk-size.
    """
    with open(file, "rb") as f:
        while True:
            chunk = f.read(chunk_size)
            if chunk:
                yield chunk
            else:
                break


def download_file_using_http_uri(http_uri, download_path, chunk_size=100000000, headers=None):
    """
    Downloads a file specified using the `http_uri` to a local `download_path`. This function
    uses a `chunk_size` to ensure an OOM error is not raised a large file is downloaded.

    Note : This function is meant to download files using presigned urls from various cloud
            providers.
    """
    if headers is None:
        headers = {}
    with cloud_storage_http_request("get", http_uri, stream=True, headers=headers) as response:
        augmented_raise_for_status(response)
        with open(download_path, "wb") as output_file:
            for chunk in response.iter_content(chunk_size=chunk_size):
                if not chunk:
                    break
                output_file.write(chunk)


@dataclass(frozen=True)
class _Chunk:
    index: int
    start: int
    end: int
    path: str


def _yield_chunks(path, file_size, chunk_size):
    num_requests = int(math.ceil(file_size / float(chunk_size)))
    for i in range(num_requests):
        range_start = i * chunk_size
        range_end = min(range_start + chunk_size - 1, file_size - 1)
        yield _Chunk(i, range_start, range_end, path)


def parallelized_download_file_using_http_uri(
    thread_pool_executor,
    http_uri,
    download_path,
    remote_file_path,
    file_size,
    uri_type,
    chunk_size,
    env,
    headers=None,
):
    """
    Downloads a file specified using the `http_uri` to a local `download_path`. This function
    sends multiple requests in parallel each specifying its own desired byte range as a header,
    then reconstructs the file from the downloaded chunks. This allows for downloads of large files
    without OOM risk.

    Note : This function is meant to download files using presigned urls from various cloud
            providers.
    Returns a dict of chunk index : exception, if one was thrown for that index.
    """

    def run_download(chunk: _Chunk):
        try:
            subprocess.run(
                [
                    sys.executable,
                    download_cloud_file_chunk.__file__,
                    "--range-start",
                    str(chunk.start),
                    "--range-end",
                    str(chunk.end),
                    "--headers",
                    json.dumps(headers or {}),
                    "--download-path",
                    download_path,
                    "--http-uri",
                    http_uri,
                ],
                text=True,
                check=True,
                capture_output=True,
                timeout=MLFLOW_DOWNLOAD_CHUNK_TIMEOUT.get(),
                env=env,
            )
        except (TimeoutExpired, CalledProcessError) as e:
            raise MlflowException(
                f"""
----- stdout -----
{e.stdout.strip()}

----- stderr -----
{e.stderr.strip()}
"""
            ) from e

    chunks = _yield_chunks(remote_file_path, file_size, chunk_size)
    # Create file if it doesn't exist or erase the contents if it does. We should do this here
    # before sending to the workers so they can each individually seek to their respective positions
    # and write chunks without overwriting.
    with open(download_path, "w"):
        pass
    if uri_type == ArtifactCredentialType.GCP_SIGNED_URL or uri_type is None:
        chunk = next(chunks)
        # GCP files could be transcoded, in which case the range header is ignored.
        # Test if this is the case by downloading one chunk and seeing if it's larger than the
        # requested size. If yes, let that be the file; if not, continue downloading more chunks.
        download_chunk(
            range_start=chunk.start,
            range_end=chunk.end,
            headers=headers,
            download_path=download_path,
            http_uri=http_uri,
        )
        downloaded_size = os.path.getsize(download_path)
        # If downloaded size was equal to the chunk size it would have been downloaded serially,
        # so we don't need to consider this here
        if downloaded_size > chunk_size:
            return {}

    futures = {thread_pool_executor.submit(run_download, chunk): chunk for chunk in chunks}
    failed_downloads = {}
    with ArtifactProgressBar.chunks(file_size, f"Downloading {download_path}", chunk_size) as pbar:
        for future in as_completed(futures):
            chunk = futures[future]
            try:
                future.result()
            except Exception as e:
                _logger.debug(
                    f"Failed to download chunk {chunk.index} for {chunk.path}: {e}. "
                    f"The download of this chunk will be retried later."
                )
                failed_downloads[chunk] = future.exception()
            else:
                pbar.update()

    return failed_downloads


def download_chunk_retries(*, chunks, http_uri, headers, download_path):
    num_retries = _MLFLOW_MPD_NUM_RETRIES.get()
    interval = _MLFLOW_MPD_RETRY_INTERVAL_SECONDS.get()
    for chunk in chunks:
        _logger.info(f"Retrying download of chunk {chunk.index} for {chunk.path}")
        for retry in range(num_retries):
            try:
                download_chunk(
                    range_start=chunk.start,
                    range_end=chunk.end,
                    headers=headers,
                    download_path=download_path,
                    http_uri=http_uri,
                )
                _logger.info(f"Successfully downloaded chunk {chunk.index} for {chunk.path}")
                break
            except Exception:
                if retry == num_retries - 1:
                    raise
            time.sleep(interval)


def _handle_readonly_on_windows(func, path, exc_info):
    """
    This function should not be called directly but should be passed to `onerror` of
    `shutil.rmtree` in order to reattempt the removal of a read-only file after making
    it writable on Windows.

    References:
    - https://bugs.python.org/issue19643
    - https://bugs.python.org/issue43657
    """
    exc_type, exc_value = exc_info[:2]
    should_reattempt = (
        os.name == "nt"
        and func in (os.unlink, os.rmdir)
        and issubclass(exc_type, PermissionError)
        and exc_value.winerror == 5
    )
    if not should_reattempt:
        raise exc_value
    os.chmod(path, stat.S_IWRITE)
    func(path)


def _get_tmp_dir():
    from mlflow.utils.databricks_utils import get_repl_id, is_in_databricks_runtime

    if is_in_databricks_runtime():
        try:
            return _get_dbutils().entry_point.getReplLocalTempDir()
        except Exception:
            pass

        if repl_id := get_repl_id():
            return os.path.join("/tmp", "repl_tmp_data", repl_id)

    return None


def create_tmp_dir():
    if directory := _get_tmp_dir():
        os.makedirs(directory, exist_ok=True)
        return tempfile.mkdtemp(dir=directory)

    return tempfile.mkdtemp()


@cache_return_value_per_process
def get_or_create_tmp_dir():
    """
    Get or create a temporary directory which will be removed once python process exit.
    """
    from mlflow.utils.databricks_utils import get_repl_id, is_in_databricks_runtime

    if is_in_databricks_runtime() and get_repl_id() is not None:
        # Note: For python process attached to databricks notebook, atexit does not work.
        # The directory returned by `dbutils.entry_point.getReplLocalTempDir()`
        # will be removed once databricks notebook detaches.
        # The temp directory is designed to be used by all kinds of applications,
        # so create a child directory "mlflow" for storing mlflow temp data.
        try:
            repl_local_tmp_dir = _get_dbutils().entry_point.getReplLocalTempDir()
        except Exception:
            repl_local_tmp_dir = os.path.join("/tmp", "repl_tmp_data", get_repl_id())

        tmp_dir = os.path.join(repl_local_tmp_dir, "mlflow")
        os.makedirs(tmp_dir, exist_ok=True)
    else:
        tmp_dir = tempfile.mkdtemp()
        # mkdtemp creates a directory with permission 0o700
        # change it to be 0o777 to ensure it can be seen in spark UDF
        os.chmod(tmp_dir, 0o777)
        atexit.register(shutil.rmtree, tmp_dir, ignore_errors=True)

    return tmp_dir


@cache_return_value_per_process
def get_or_create_nfs_tmp_dir():
    """
    Get or create a temporary NFS directory which will be removed once python process exit.
    """
    from mlflow.utils.databricks_utils import get_repl_id, is_in_databricks_runtime
    from mlflow.utils.nfs_on_spark import get_nfs_cache_root_dir

    nfs_root_dir = get_nfs_cache_root_dir()

    if is_in_databricks_runtime() and get_repl_id() is not None:
        # Note: In databricks, atexit hook does not work.
        # The directory returned by `dbutils.entry_point.getReplNFSTempDir()`
        # will be removed once databricks notebook detaches.
        # The temp directory is designed to be used by all kinds of applications,
        # so create a child directory "mlflow" for storing mlflow temp data.
        try:
            repl_nfs_tmp_dir = _get_dbutils().entry_point.getReplNFSTempDir()
        except Exception:
            repl_nfs_tmp_dir = os.path.join(nfs_root_dir, "repl_tmp_data", get_repl_id())

        tmp_nfs_dir = os.path.join(repl_nfs_tmp_dir, "mlflow")
        os.makedirs(tmp_nfs_dir, exist_ok=True)
    else:
        tmp_nfs_dir = tempfile.mkdtemp(dir=nfs_root_dir)
        # mkdtemp creates a directory with permission 0o700
        # change it to be 0o777 to ensure it can be seen in spark UDF
        os.chmod(tmp_nfs_dir, 0o777)
        atexit.register(shutil.rmtree, tmp_nfs_dir, ignore_errors=True)

    return tmp_nfs_dir


def write_spark_dataframe_to_parquet_on_local_disk(spark_df, output_path):
    """
    Write spark dataframe in parquet format to local disk.

    :param spark_df: Spark dataframe
    :param output_path: path to write the data to
    """
    from mlflow.utils.databricks_utils import is_in_databricks_runtime

    if is_in_databricks_runtime():
        dbfs_path = os.path.join(".mlflow", "cache", str(uuid.uuid4()))
        spark_df.coalesce(1).write.format("parquet").save(dbfs_path)
        shutil.copytree("/dbfs/" + dbfs_path, output_path)
        shutil.rmtree("/dbfs/" + dbfs_path)
    else:
        spark_df.coalesce(1).write.format("parquet").save(output_path)


def shutil_copytree_without_file_permissions(src_dir, dst_dir):
    """
    Copies the directory src_dir into dst_dir, without preserving filesystem permissions
    """
    for dirpath, dirnames, filenames in os.walk(src_dir):
        for dirname in dirnames:
            relative_dir_path = os.path.relpath(os.path.join(dirpath, dirname), src_dir)
            # For each directory <dirname> immediately under <dirpath>, create an equivalently-named
            # directory under the destination directory
            abs_dir_path = os.path.join(dst_dir, relative_dir_path)
            os.mkdir(abs_dir_path)
        for filename in filenames:
            # For each file with name <filename> immediately under <dirpath>, copy that file to
            # the appropriate location in the destination directory
            file_path = os.path.join(dirpath, filename)
            relative_file_path = os.path.relpath(file_path, src_dir)
            abs_file_path = os.path.join(dst_dir, relative_file_path)
            shutil.copy2(file_path, abs_file_path)


def contains_path_separator(path):
    """
    Returns True if a path contains a path separator, False otherwise.
    """
    return any((sep in path) for sep in (os.path.sep, os.path.altsep) if sep is not None)


def read_chunk(path: os.PathLike, size: int, start_byte: int = 0) -> bytes:
    """
    Read a chunk of bytes from a file.

    :param path: Path to the file.
    :param size: The size of the chunk.
    :param start_byte: The start byte of the chunk.
    :return: The chunk of bytes.
    """
    with open(path, "rb") as f:
        if start_byte > 0:
            f.seek(start_byte)
        return f.read(size)


@contextmanager
def remove_on_error(path: os.PathLike, onerror=None):
    """
    A context manager that removes a file or directory if an exception is raised during execution.

    :param path: Path to the file or directory.
    :param onerror: A callback function that will be called with the captured exception before
                    the file or directory is removed. For example, you can use this callback to
                    log the exception.
    """
    try:
        yield
    except Exception as e:
        if onerror:
            onerror(e)
        if os.path.exists(path):
            if os.path.isfile(path):
                os.remove(path)
            elif os.path.isdir(path):
                shutil.rmtree(path)
        raise


@contextmanager
def chdir(path: str) -> None:
    """
    Temporarily change the current working directory to the specified path.

    :param path: The path to use as the temporary working directory.
    """
    cwd = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(cwd)


def get_total_file_size(path: Union[str, pathlib.Path]) -> Optional[int]:
    """
    Return the size of all files under given path, including files in subdirectories.

    :param path: The absolute path of a local directory.
    :return: size in bytes.
    """
    try:
        if isinstance(path, pathlib.Path):
            path = str(path)
        if not os.path.exists(path):
            raise MlflowException(
                message=f"The given {path} does not exist.", error_code=INVALID_PARAMETER_VALUE
            )
        if not os.path.isdir(path):
            raise MlflowException(
                message=f"The given {path} is not a directory.", error_code=INVALID_PARAMETER_VALUE
            )

        total_size = 0
        for cur_path, dirs, files in os.walk(path):
            full_paths = [os.path.join(cur_path, file) for file in files]
            total_size += sum([os.path.getsize(file) for file in full_paths])
        return total_size
    except Exception as e:
        _logger.info(f"Failed to get the total size of {path} because of error :{e}")
        return None
