# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT license.

import os
import subprocess

import numpy as np


def output_vectors(
    diskann_build_path: str, temporary_file_path: str, vectors: np.ndarray, timeout: int = 60
) -> str:
    vectors_as_tsv_path = os.path.join(temporary_file_path, "vectors.tsv")
    with open(vectors_as_tsv_path, "w") as vectors_tsv_out:
        for vector in vectors:
            as_str = "\t".join(str(component) for component in vector)
            print(as_str, file=vectors_tsv_out)
    # there is probably a clever way to have numpy write out C++ friendly floats, so feel free to remove this in
    # favor of something more sane later
    vectors_as_bin_path = os.path.join(temporary_file_path, "vectors.bin")
    tsv_to_bin_path = os.path.join(diskann_build_path, "apps", "utils", "tsv_to_bin")

    number_of_points, dimensions = vectors.shape
    args = [
        tsv_to_bin_path,
        "float",
        vectors_as_tsv_path,
        vectors_as_bin_path,
        str(dimensions),
        str(number_of_points),
    ]
    completed = subprocess.run(args, timeout=timeout)
    if completed.returncode != 0:
        raise Exception(
            f"Unable to convert tsv to binary using tsv_to_bin, completed_process: {completed}"
        )
    return vectors_as_bin_path


def build_ssd_index(
    diskann_build_path: str,
    temporary_file_path: str,
    vectors: np.ndarray,
    per_process_timeout: int = 60,  # this may not be long enough if you're doing something larger
):
    vectors_as_bin_path = output_vectors(
        diskann_build_path, temporary_file_path, vectors, timeout=per_process_timeout
    )

    ssd_builder_path = os.path.join(diskann_build_path, "apps", "build_disk_index")
    args = [
        ssd_builder_path,
        "--data_type",
        "float",
        "--dist_fn",
        "l2",
        "--data_path",
        vectors_as_bin_path,
        "--index_path_prefix",
        os.path.join(temporary_file_path, "smoke_test"),
        "-R",
        "64",
        "-L",
        "100",
        "--search_DRAM_budget",
        "1",
        "--build_DRAM_budget",
        "1",
        "--num_threads",
        "1",
        "--PQ_disk_bytes",
        "0",
    ]
    completed = subprocess.run(args, timeout=per_process_timeout)

    if completed.returncode != 0:
        command_run = " ".join(args)
        raise Exception(
            f"Unable to build a disk index with the command: '{command_run}'\ncompleted_process: {completed}\nstdout: {completed.stdout}\nstderr: {completed.stderr}"
        )
    # index is now built inside of temporary_file_path
