# Copyright (c) 2022-2023, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
from pathlib import Path

from cugraph_service_server import (
    defaults,
    start_server_blocking,
)


def main():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "--host",
        type=str,
        default=defaults.host,
        help="hostname the server should use, default is " f"{defaults.host}",
    )
    arg_parser.add_argument(
        "--port",
        type=int,
        default=defaults.port,
        help="port the server should listen on, default " f"is {defaults.port}",
    )
    arg_parser.add_argument(
        "--graph-creation-extension-dir",
        type=Path,
        help="dir to load graph creation extension " "functions from",
    )
    arg_parser.add_argument(
        "--dask-scheduler-file",
        type=Path,
        help="file generated by a dask scheduler, used "
        "for connecting to a dask cluster for MG support",
    )
    arg_parser.add_argument(
        "--start-local-cuda-cluster",
        action="store_true",
        help="use a LocalCUDACluster for multi-GPU",
    )
    arg_parser.add_argument(
        "--protocol",
        type=str,
        help="protocol to use by the dask cluster, ex 'ucx'",
    )
    arg_parser.add_argument(
        "--rmm-pool-size",
        type=str,
        help="RMM pool size, ex. '28GB'",
    )
    arg_parser.add_argument(
        "--dask-worker-devices",
        type=str,
        help="list of GPU device IDs the dask cluster should use, ex. '0,1,2,3'",
    )
    args = arg_parser.parse_args()

    msg = "Starting the cugraph_service server "
    if (args.dask_scheduler_file is None) and (args.start_local_cuda_cluster is False):
        msg += "(single-GPU)..."
    elif (args.dask_scheduler_file is not None) and (
        args.start_local_cuda_cluster is True
    ):
        raise RuntimeError(
            "dask-scheduler-file cannot be set if start-local-cuda-cluster is specified"
        )
    elif args.dask_scheduler_file is not None:
        msg += f"(multi-GPU, scheduler file: {args.dask_scheduler_file})..."
    elif args.start_local_cuda_cluster is True:
        msg += "(multi-GPU, LocalCUDACluster)..."

    start_server_blocking(
        args.graph_creation_extension_dir,
        args.start_local_cuda_cluster,
        args.dask_scheduler_file,
        args.host,
        args.port,
        args.protocol,
        args.rmm_pool_size,
        args.dask_worker_devices,
        console_message=msg,
    )
    print("done.")


if __name__ == "__main__":
    main()
