import sys
import subprocess
import os
import tempfile
from types import ModuleType
from typing import Dict, List

base_install_dir = tempfile.gettempdir() + "/persistence_test_chromadb_versions"


def get_path_to_version_install(version: str) -> str:
    return base_install_dir + "/" + version


def switch_to_version(version: str, versioned_modules: List[str]) -> ModuleType:
    module_name = "chromadb"
    # Remove old version from sys.modules, except test modules
    old_modules = {
        n: m
        for n, m in sys.modules.items()
        if n == module_name
        or (n.startswith(module_name + "."))
        or n in versioned_modules
        or (any(n.startswith(m + ".") for m in versioned_modules))
    }
    for n in old_modules:
        del sys.modules[n]
    # Load the target version and override the path to the installed version
    # https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
    sys.path.insert(0, get_path_to_version_install(version))
    import chromadb

    assert chromadb.__version__ == version
    return chromadb


def get_path_to_version_library(version: str) -> str:
    return get_path_to_version_install(version) + "/chromadb/__init__.py"


def install_version(version: str, dep_overrides: Dict[str, str]) -> None:
    # Check if already installed
    version_library = get_path_to_version_library(version)
    if os.path.exists(version_library):
        return
    path = get_path_to_version_install(version)
    install(f"chromadb=={version}", path, dep_overrides)


def install(pkg: str, path: str, dep_overrides: Dict[str, str]) -> int:
    # -q -q to suppress pip output to ERROR level
    # https://pip.pypa.io/en/stable/cli/pip/#quiet
    print("Purging pip cache")
    subprocess.check_call(
        [
            sys.executable,
            "-m",
            "pip",
            "cache",
            "purge",
        ]
    )

    command = [sys.executable, "-m", "pip", "-q", "-q", "install", pkg]

    for dep, operator_version in dep_overrides.items():
        command.append(f"{dep}{operator_version}")

    command.append("--no-binary=chroma-hnswlib")
    command.append(f"--target={path}")

    print(f"Installing chromadb version {pkg} to {path}")
    return subprocess.check_call(command)