python3 << EndPython3
import collections
import os
import sys
import vim


def strtobool(text):
    if text.lower() in ["y", "yes", "t", "true", "on", "1"]:
        return True
    if text.lower() in ["n", "no", "f", "false", "off", "0"]:
        return False
    raise ValueError(f"{text} is not convertible to boolean")


class Flag(collections.namedtuple("FlagBase", "name, cast")):
    @property
    def var_name(self):
        return self.name.replace("-", "_")

    @property
    def vim_rc_name(self):
        name = self.var_name
        if name == "line_length":
            name = name.replace("_", "")
        return "g:black_" + name


FLAGS = [
    Flag(name="line_length", cast=int),
    Flag(name="fast", cast=strtobool),
    Flag(name="skip_string_normalization", cast=strtobool),
    Flag(name="quiet", cast=strtobool),
    Flag(name="skip_magic_trailing_comma", cast=strtobool),
    Flag(name="preview", cast=strtobool),
]


def _get_python_binary(exec_prefix, pyver):
    try:
        default = vim.eval("g:pymode_python").strip()
    except vim.error:
        default = ""
    if default and os.path.exists(default):
        return default
    if sys.platform[:3] == "win":
        return exec_prefix / "python.exe"
    bin_path = exec_prefix / "bin"
    exec_path = (bin_path / f"python{pyver[0]}.{pyver[1]}").resolve()
    if exec_path.exists():
        return exec_path
    # It is possible that some environments may only have python3
    exec_path = (bin_path / "python3").resolve()
    if exec_path.exists():
        return exec_path
    raise ValueError("python executable not found")


def _get_pip(venv_path):
    if sys.platform[:3] == "win":
        return venv_path / "Scripts" / "pip.exe"
    return venv_path / "bin" / "pip"


def _get_virtualenv_site_packages(venv_path, pyver):
    if sys.platform[:3] == "win":
        return venv_path / "Lib" / "site-packages"
    if (
        venv_path.exists()
        and not (venv_path / "lib" / f"python{pyver[0]}.{pyver[1]}").exists()
    ):
        # The virtualenv already exists but it doesn't seem to have the expected
        # Python version, so we disregard the requested `pyver` and
        # discover the real Python interpreter from this virtualenv
        import subprocess

        result = subprocess.run(
            [_get_python_binary(venv_path, pyver), "--version"],
            stdout=subprocess.PIPE,
            text=True,
        )
        venv_version = result.stdout.split(" ")[1].strip().split(".")
        return (
            venv_path
            / "lib"
            / f"python{venv_version[0]}.{venv_version[1]}"
            / "site-packages"
        )
    else:
        return venv_path / "lib" / f"python{pyver[0]}.{pyver[1]}" / "site-packages"


def _initialize_black_env(upgrade=False):
    if vim.eval("g:black_use_virtualenv ? 'true' : 'false'") == "false":
        if upgrade:
            print("Upgrade disabled due to g:black_use_virtualenv being disabled.")
            print(
                "Either use your system package manager (or pip) to upgrade black"
                " separately,"
            )
            print("or modify your vimrc to have 'let g:black_use_virtualenv = 1'.")
            return False
        else:
            # Nothing needed to be done.
            return True

    pyver = sys.version_info[:3]
    if pyver < (3, 10):
        print("Sorry, Black requires Python 3.10+ to run.")
        return False

    from pathlib import Path
    import subprocess
    import venv

    virtualenv_path = Path(vim.eval("g:black_virtualenv")).expanduser()
    virtualenv_site_packages = str(
        _get_virtualenv_site_packages(virtualenv_path, pyver)
    )
    first_install = False
    if not virtualenv_path.is_dir():
        print("Please wait, one time setup for Black.")
        _executable = sys.executable
        _base_executable = getattr(sys, "_base_executable", _executable)
        try:
            executable = str(_get_python_binary(Path(sys.exec_prefix), pyver))
            sys.executable = executable
            sys._base_executable = executable
            print(f"Creating a virtualenv in {virtualenv_path}...")
            print(
                "(this path can be customized in .vimrc by setting g:black_virtualenv)"
            )
            venv.create(virtualenv_path, with_pip=True)
        except Exception:
            print(
                "Encountered exception while creating virtualenv (see traceback below)."
            )
            print(f"Removing {virtualenv_path}...")
            import shutil

            shutil.rmtree(virtualenv_path)
            raise
        finally:
            sys.executable = _executable
            sys._base_executable = _base_executable
        first_install = True
    if first_install:
        print("Installing Black with pip...")
    if upgrade:
        print("Upgrading Black with pip...")
    if first_install or upgrade:
        subprocess.run(
            [str(_get_pip(virtualenv_path)), "install", "-U", "black"],
            stdout=subprocess.PIPE,
        )
        print("DONE! You are all set, thanks for waiting ✨ 🍰 ✨")
    if first_install:
        print(
            "Pro-tip: to upgrade Black in the future, use the :BlackUpgrade command and"
            " restart Vim.\n"
        )
    if virtualenv_site_packages not in sys.path:
        sys.path.insert(0, virtualenv_site_packages)
    return True


if _initialize_black_env():
    try:
        import black
    except ImportError:
        print(f"Could not import black from any of: {', '.join(sys.path)}.")
    import time


def get_target_version(tv):
    if isinstance(tv, black.TargetVersion):
        return tv
    ret = None
    try:
        ret = black.TargetVersion[tv.upper()]
    except KeyError:
        print(
            f"WARNING: Target version {tv!r} not recognized by Black, using default"
            " target"
        )
    return ret


def Black(**kwargs):
    """
    kwargs allows you to override ``target_versions`` argument of
    ``black.FileMode``.

    ``target_version`` needs to be cleaned because ``black.FileMode``
    expects the ``target_versions`` argument to be a set of TargetVersion enums.

    Allow kwargs["target_version"] to be a string to allow
    to type it more quickly.

    Using also target_version instead of target_versions to remain
    consistent to Black's documentation of the structure of pyproject.toml.
    """
    start = time.time()
    configs = get_configs()

    black_kwargs = {}
    if "target_version" in kwargs:
        target_version = kwargs["target_version"]

        if not isinstance(target_version, (list, set)):
            target_version = [target_version]
        target_version = set(
            filter(lambda x: x, map(lambda tv: get_target_version(tv), target_version))
        )
        black_kwargs["target_versions"] = target_version

    mode = black.FileMode(
        line_length=configs["line_length"],
        string_normalization=not configs["skip_string_normalization"],
        is_pyi=vim.current.buffer.name.endswith(".pyi"),
        magic_trailing_comma=not configs["skip_magic_trailing_comma"],
        preview=configs["preview"],
        **black_kwargs,
    )
    quiet = configs["quiet"]

    buffer_str = "\n".join(vim.current.buffer) + "\n"
    try:
        new_buffer_str = black.format_file_contents(
            buffer_str,
            fast=configs["fast"],
            mode=mode,
        )
    except black.NothingChanged:
        if not quiet:
            print(
                "Black: already well formatted, good job. (took"
                f" {time.time() - start:.4f}s)"
            )
    except Exception as exc:
        print(f"Black: {exc}")
    else:
        current_buffer = vim.current.window.buffer
        cursors = []
        for i, tabpage in enumerate(vim.tabpages):
            if tabpage.valid:
                for j, window in enumerate(tabpage.windows):
                    if window.valid and window.buffer == current_buffer:
                        cursors.append((i, j, window.cursor))
        vim.current.buffer[:] = new_buffer_str.split("\n")[:-1]
        for i, j, cursor in cursors:
            window = vim.tabpages[i].windows[j]
            try:
                window.cursor = cursor
            except vim.error:
                window.cursor = (len(window.buffer), 0)
        if not quiet:
            print(f"Black: reformatted in {time.time() - start:.4f}s.")


def get_configs():
    filename = vim.eval("@%")
    path_pyproject_toml = black.find_pyproject_toml((filename,))
    if path_pyproject_toml:
        toml_config = black.parse_pyproject_toml(path_pyproject_toml)
    else:
        toml_config = {}

    return {
        flag.var_name: toml_config.get(flag.name, flag.cast(vim.eval(flag.vim_rc_name)))
        for flag in FLAGS
    }


def BlackUpgrade():
    _initialize_black_env(upgrade=True)


def BlackVersion():
    print(f"Black, version {black.__version__} on Python {sys.version}.")
EndPython3

function black#Black(...)
    let kwargs = {}
    for arg in a:000
        let arg_list = split(arg, '=')
        let kwargs[arg_list[0]] = arg_list[1]
    endfor
python3 << EOF
import vim

kwargs = vim.eval("kwargs")
EOF
  :py3 Black(**kwargs)
endfunction

function black#BlackUpgrade()
  :py3 BlackUpgrade()
endfunction

function black#BlackVersion()
  :py3 BlackVersion()
endfunction