# SPDX-FileCopyrightText: 2026 MetaMachines LLC
#
# SPDX-License-Identifier: MIT

from cuda.core import Device, Program, ProgramOptions

from mm_ptx import get_include_dir

import sys


class NvCompilerHelper:
    def __init__(self):
        self.dev = Device()
        self.dev.set_current()
        capability = self.dev.compute_capability
        self.arch = f"sm_{capability.major}{capability.minor}"
        self.include_dirs = [get_include_dir()]

    def _make_program_options(self, **kwargs):
        include_dirs = kwargs.pop("include_dirs", None)
        if include_dirs:
            for include_key in ("include_path", "include_paths", "include_dirs"):
                try:
                    return ProgramOptions(**kwargs, **{include_key: include_dirs})
                except TypeError:
                    continue
        return ProgramOptions(**kwargs)

    def cuda_to_ptx(self, cuda: str):
        program_options = self._make_program_options(std="c++11", arch=self.arch, include_dirs=self.include_dirs)
        prog = Program(cuda, code_type="c++", options=program_options)
        mod = prog.compile("ptx", logs=sys.stdout)
        return mod.code.decode("utf-8")

    def ptx_to_cubin(self, ptx: str):
        program_options = ProgramOptions(arch=self.arch)
        prog = Program(ptx, code_type="ptx", options=program_options)

        mod = prog.compile("cubin", logs=sys.stdout)
        return mod
