# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import click
import jinja2
import os


@click.command()
@click.option("--nvshmem-home", type=click.Path(exists=True), required=True)
@click.option("--config-version", type=str, required=True)
@click.option("--binding-name", type=str, required=True)
@click.option("--entry-point-path", type=click.Path(exists=True), required=True)
@click.option("--input-path", type=click.Path(exists=True), required=True)
@click.option("--output-path", type=click.Path(), required=True)
def main(nvshmem_home, config_version, binding_name, entry_point_path, input_path, output_path):
    # Load the Jinja2 template
    template_dir = os.path.dirname(input_path)
    template_name = os.path.basename(input_path)

    env = jinja2.Environment(loader=jinja2.FileSystemLoader(template_dir))
    template = env.get_template(template_name)

    # Get CUDA13 CCCL include path
    CUDA_HOME = os.environ.get("CUDA_HOME", "/usr/local/cuda")
    CUDA13_CCCL_INCLUDE_PATH = os.path.join(CUDA_HOME, "include", "cccl")

    if not os.path.exists(CUDA13_CCCL_INCLUDE_PATH):
        CUDA13_CCCL_INCLUDE_PATH = ""
    else:
        CUDA13_CCCL_INCLUDE_PATH = f"- {CUDA13_CCCL_INCLUDE_PATH}"

    # Prepare template variables
    template_vars = {
        'CONFIG_VERSION': config_version,
        'NVSHMEM_HOME': nvshmem_home,
        'CUDA13_CCCL_INCLUDE_PATH': CUDA13_CCCL_INCLUDE_PATH,
        'ENTRY_POINT_PATH': entry_point_path,
        'OUTPUT_NAME': binding_name
    }

    # Render the template
    rendered_content = template.render(**template_vars)

    # Create output directory if it doesn't exist
    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Write the rendered content to the output file
    with open(output_path, 'w') as f:
        f.write(rendered_content)

    print(f"Successfully generated YAML config file: {output_path}")


if __name__ == "__main__":
    main()
