"""Build a private target-specific texture override from the user's local DLSS DLL.

The resulting PTX/code object contains NVIDIA code and must remain local.
No Wine, shell script, GPU, CUDA SDK or system installation is required.
"""
import argparse
import ast
import hashlib
import json
import os
import pathlib
import re
import shutil
import subprocess
import sys
from gpu_targets import TARGETS
from llvm22_ir import compatible_ir


def require_code_object_target(path, target):
    with path.open('rb') as file:
        header = file.read(64)
    if (len(header) != 64 or header[:6] != b'\x7fELF\x02\x01' or
            header[18:20] != b'\xe0\x00' or header[48] != TARGETS[target]['elfMachine']):
        raise RuntimeError(f'Code object does not target {target}: {path}')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--dlss-dll', type=pathlib.Path, required=True)
    parser.add_argument('--hip-root', type=pathlib.Path, required=True)
    parser.add_argument('--zluda-root', type=pathlib.Path, required=True)
    parser.add_argument('--kernel', required=True)
    parser.add_argument('--output-directory', type=pathlib.Path, required=True)
    parser.add_argument('--gpu-arch', choices=tuple(TARGETS), default='gfx1201')
    parser.add_argument('--shader-mode', choices=['cu', 'wgp'], default='cu')
    parser.add_argument('--wave64', action='store_true', help='experimental wave64 for M post kernels without matrix operations')
    args = parser.parse_args()
    is_m_post = bool(re.fullmatch(r'rrlite_post_3_[12]_mv(?:hi|lo)_(?:hdr|ldr)_folded', args.kernel))
    if not is_m_post and not re.fullmatch(r'hiluma_engine_(?:output_depth(?:inv|reg)_mv(?:hi|lo)_(?:hdr|ldr)(?:_max)?|'
                                        r'input_depth(?:inv|reg)_mv(?:hi|lo)_(?:hdr|ldr))_v2_rel', args.kernel):
        parser.error('this private experimental build recipe supports K v2 input/output and M folded post variants only')
    if args.wave64 and not is_m_post:
        parser.error('--wave64 is restricted to M post kernels without matrix operations')
    repo = pathlib.Path(__file__).resolve().parents[2]
    output = args.output_directory.resolve()
    if not output.is_relative_to((repo / 'build').resolve()):
        parser.error('private texture outputs must stay under the repository ignored build directory')
    output.mkdir(parents=True, exist_ok=True)
    for existing in output.glob('*.hsaco'):
        require_code_object_target(existing, args.gpu_arch)
    work = output / ('work-' + args.kernel)
    work.mkdir(exist_ok=True)
    hip = args.hip_root.resolve()
    zluda = args.zluda_root.resolve()
    tools = {}
    for name in ('clang++', 'llvm-dis', 'llvm-as'):
        candidates = [hip / folder / (name + '.exe') for folder in ('lib/llvm/bin', 'bin')]
        tools[name] = next((path for path in candidates if path.is_file()), candidates[0])
    emitter = zluda / 'd4r_emit.exe'
    # The emitter's LLVM reader must consume the helper bitcode. A newer HIP
    # compiler can emit textual IR, but its bitcode attributes may be unreadable.
    emitter_assembler = zluda / 'llvm-as.exe'
    if emitter_assembler.is_file():
        tools['llvm-as'] = emitter_assembler
    for path in [args.dlss_dll, emitter, *tools.values()]:
        if not path.is_file():
            parser.error(f'missing dependency: {path}')
    env = os.environ.copy()
    env.update(D4R_PREFER_ACCURACY='1', D4R_ZLUDA_WMMA='1', D4R_ZLUDA_WMMA_FP8='1',
               D4R_ZLUDA_WMMA_FP8_NATIVE='0', D4R_ZLUDA_WMMA_F16_REFERENCE='1',
               D4R_ZLUDA_WGP='1' if args.shader_mode == 'wgp' else '0',
               D4R_ZLUDA_IMPLICIT_MAX_BLOCK='256')
    for name in ('D4R_ZLUDA_IGNORE_DENORMAL', 'D4R_ZLUDA_FAST_MATH', 'D4R_ZLUDA_WMMA_F32ACC',
                 'D4R_ZLUDA_WAVE64', 'D4R_ZLUDA_PROFILE'):
        env.pop(name, None)
    if args.wave64:
        env['D4R_ZLUDA_WAVE64'] = '1'
    env['PATH'] = str(zluda) + os.pathsep + str(hip / 'bin') + os.pathsep + env['PATH']

    def run(label, command):
        with (work / (label + '.log')).open('wb') as log:
            result = subprocess.run([str(arg) for arg in command], env=env, stdout=log,
                                    stderr=subprocess.STDOUT, timeout=600)
        if result.returncode:
            print((work / (label + '.log')).read_text(errors='replace')[-8000:], file=sys.stderr)
            raise RuntimeError(f'{label} failed ({result.returncode}); see {work}')

    ptx = work / 'ptx'
    run('extract', [sys.executable, repo / 'kernels/tools/extract_dlss_ptx.py', args.dlss_dll, ptx])
    env['D4R_DLSS_PTX_DIR'] = str(ptx)
    edited = work / (args.kernel + '.ptx')
    run('make-ptx', [sys.executable, repo / 'kernels/tex/make_ptx.py', args.kernel, edited])
    rewrite_match = re.search(r'surface stores replaced: (\{[^\n]+\})',
                              (work / 'make-ptx.log').read_text(errors='replace'))
    if not rewrite_match:
        raise RuntimeError('PTX recipe did not report its changes')
    rewrite_counts = ast.literal_eval(rewrite_match[1])
    if not any(rewrite_counts.values()):
        raise RuntimeError('PTX recipe made no changes; refuse a no-op override')
    raw, ir, extra = work / 'raw.bc', work / 'raw.ll', work / 'extra.bc'
    run('compile-helper', [tools['clang++'], '-x', 'hip', '-std=c++20', '-nogpuinc', '-nogpulib', '-O3',
                          '-mwavefrontsize64' if args.wave64 else '-mno-wavefrontsize64',
                          '--offload-device-only', '--offload-arch=' + args.gpu_arch,
                          '-fgpu-rdc', '-emit-llvm', '-c', '-Xclang', '-fdenormal-fp-math=dynamic',
                          '-DD4R_ACCURACY', '-DD4R_KERNEL_NAME=' + args.kernel,
                          '-o', raw, repo / 'kernels/tex/sust_only.hip'])
    run('disassemble', [tools['llvm-dis'], raw, '-o', ir])
    lines = []
    for line in ir.read_text().splitlines():
        # Helpers are linked directly into the translated module. Their ThinLTO
        # index becomes stale when the unused HIP/LLVM globals below are removed.
        if re.match(r'^\^\d+\s*=', line):
            continue
        if re.search(r'@llvm.used|wchar_size|llvm.module.flags|__hip_cuid', line):
            continue
        line = line.replace('optnone', '').replace(f'"target-cpu"="{args.gpu_arch}"', '')
        lines.append(re.sub(r'"target-features"="[^"]+"', '', line))
    ir.write_text(compatible_ir('\n'.join(lines) + '\n') if emitter_assembler.is_file()
                  else '\n'.join(lines) + '\n')
    run('assemble', [tools['llvm-as'], ir, '-o', extra])
    env['D4R_ZLUDA_EXTRA_BC'] = str(extra)
    run('emit', [emitter, edited, work / 'emitted', args.gpu_arch])
    optimized = (work / 'emitted/opt.ll').read_text()
    definition = re.search(r'^define amdgpu_kernel .*@' + re.escape(args.kernel) + r'\(.*#(\d+)\s*\{', optimized, re.M)
    if not definition:
        raise RuntimeError('Cannot verify shader mode: kernel attributes are missing')
    attributes = re.search(r'^attributes #' + definition[1] + r' = \{([^\n]+)\}', optimized, re.M)
    expected_mode = '-cumode' if args.shader_mode == 'wgp' else '+cumode'
    if not attributes or expected_mode not in attributes[1]:
        raise RuntimeError(f'Emitter did not honor {args.shader_mode} mode; build the corresponding ZLUDA patch first')
    if args.wave64 and '+wavefrontsize64' not in attributes[1]:
        raise RuntimeError('Emitter did not honor wave64; candidate is not accepted')
    require_code_object_target(work / 'emitted/module.hsaco', args.gpu_arch)
    shutil.copyfile(work / 'emitted/module.hsaco', output / (args.kernel + '.hsaco'))
    run('manifest', [sys.executable, repo / 'kernels/tools/kernel_manifest.py', output, args.dlss_dll])
    metadata = dict(architecture=args.gpu_arch, accuracy=True, nativeFP8=False, shaderMode=args.shader_mode,
                    wave64=args.wave64,
                    lastBuiltKernel=args.kernel, lastRecipeRewrites=rewrite_counts,
                    privateNvidiaDerivedCode=True, validation='not yet validated; do not install before comparison',
                    dlssSha256=hashlib.sha256(args.dlss_dll.read_bytes()).hexdigest(),
                    zludaBuild=json.loads((zluda / 'build-info.json').read_text(encoding='utf-8-sig')),
                    compiler=subprocess.check_output([str(tools['clang++']), '--version'], env=env).decode(errors='replace'),
                    objects={p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in output.glob('*.hsaco')})
    (output / 'build-info.json').write_text(json.dumps(metadata, indent=2) + '\n')
    print(f'Built PRIVATE {args.gpu_arch} texture override: {output}; validation required before use')


if __name__ == '__main__':
    main()