# Copyright (c) 2026 Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

"""Generate runtime ISA properties from the amdisa profiles."""

from pathlib import Path

_AMDGPU_ARCH_ORDER = (
    'cdna1',
    'cdna2',
    'cdna3',
    'cdna4',
    'rdna1',
    'rdna2',
    'rdna3',
    'rdna3_5',
    'rdna4',
    'gfx1250',
)


def emit_isa_properties(output_dir: str, specs) -> Path:
    """Emit a constexpr runtime property map for the supplied AMDGPU ISAs."""
    profiles = {name: spec.profile for name, spec, _ in specs}
    unknown = profiles.keys() - set(_AMDGPU_ARCH_ORDER)
    if unknown:
        raise ValueError(f'unsupported ISA property entries: {sorted(unknown)}')

    cases = []
    max_addressable_vgprs_per_wf = 0
    for name in _AMDGPU_ARCH_ORDER:
        profile = profiles.get(name)
        if profile is None:
            continue
        enum_name = name.upper()
        supports_wgp_mode = 'true' if profile.supports_wgp_mode else 'false'
        descriptor_sgpr_count_encoded = (
            'true' if profile.descriptor_sgpr_count_encoded else 'false'
        )
        uses_ttmp_workgroup_ids = 'true' if profile.uses_ttmp_workgroup_ids else 'false'
        uses_cluster_ttmp_workgroup_ids = (
            'true' if profile.uses_cluster_ttmp_workgroup_ids else 'false'
        )
        addressable_vgprs = profile.max_addressable_vgprs_per_wf
        max_addressable_vgprs_per_wf = max(
            max_addressable_vgprs_per_wf, addressable_vgprs
        )
        cases += [
            f'  case ROCJITSU_CODE_ARCH_{enum_name}:',
            '    return {',
            f'        .supports_wgp_mode = {supports_wgp_mode},',
            f'        .descriptor_sgpr_count_encoded = {descriptor_sgpr_count_encoded},',
            f'        .uses_ttmp_workgroup_ids = {uses_ttmp_workgroup_ids},',
            f'        .uses_cluster_ttmp_workgroup_ids = {uses_cluster_ttmp_workgroup_ids},',
            f'        .max_addressable_vgprs_per_wf = {addressable_vgprs},',
            '    };',
        ]

    lines = [
        '// Copyright (c) 2026 Advanced Micro Devices, Inc.',
        '// SPDX-License-Identifier: MIT',
        '//',
        '// AUTO-GENERATED by the amdisa codegen pipeline. DO NOT EDIT.',
        '// See lib/python/amdisa/README.md for regeneration instructions.',
        '',
        '#ifndef ROCJITSU_ISA_ARCH_AMDGPU_ISA_PROPERTIES_H_',
        '#define ROCJITSU_ISA_ARCH_AMDGPU_ISA_PROPERTIES_H_',
        '',
        '#include "rocjitsu/base/api.h"',
        '',
        '#include <cstdint>',
        '',
        'namespace rocjitsu {',
        '',
        'struct IsaProperties {',
        '  bool supports_wgp_mode = false;',
        '  bool descriptor_sgpr_count_encoded = true;',
        '  bool uses_ttmp_workgroup_ids = false;',
        '  bool uses_cluster_ttmp_workgroup_ids = false;',
        '  uint32_t max_addressable_vgprs_per_wf = 0;',
        '};',
        '',
        'inline constexpr uint32_t MAX_SUPPORTED_ADDRESSABLE_VGPRS_PER_WF = '
        f'{max_addressable_vgprs_per_wf};',
        '',
        '[[nodiscard]] constexpr IsaProperties isa_properties(rj_code_arch_t arch) {',
        '  switch (arch) {',
        *cases,
        '  default:',
        '    return {};',
        '  }',
        '}',
        '',
        '} // namespace rocjitsu',
        '',
        '#endif // ROCJITSU_ISA_ARCH_AMDGPU_ISA_PROPERTIES_H_',
        '',
    ]

    output = Path(output_dir) / 'isa_properties.h'
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text('\n'.join(lines))
    return output