#!/usr/bin/env python3
import argparse
from itertools import takewhile
import json
import os
import sys
from typing import Callable, Dict, IO, List
import io
from string import Template
import re

from spirv_common import ignore_overload, emit_guards, close_guards, clang_format

def ignore_function(fun):
    blacklist = [
        "GenericCastToPtrExplicit",
        "GenericPtrMemSemantics",
        "GroupAll",
        "GroupAny",
        "GroupBroadcast",
        "GroupFAdd",
        "GroupFMax",
        "GroupFMin",
        "GroupIAdd",
        "GroupSMax",
        "GroupSMin",
        "GroupUMax",
        "GroupUMin",
        "printf"
    ]
    return any([fun.find(b) != -1 for b in blacklist])

def get_builtin_name(func):
    func = func.replace("__spirv_", "")
    builtin = func.replace("ocl_", "")
    kind = "core" if builtin == func else "ocl"
    return (kind, builtin)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="""
Generate SPIR-V interface check tests.
Typical usage:
 clang-tblgen --gen-clang-progmodel-builtins-as-json path/to/SPIRVBuiltins.td \
   -o builtin.json
 generate-test-libclc.py builtin.json -format clang-format -o spirv_builtin.h
""")
    parser.add_argument("input",
                        metavar="MAPPING",
                        type=argparse.FileType('r'),
                        help="Path to mapping.json")
    parser.add_argument("-format",
                        metavar="clang-format",
                        nargs='?',
                        help="clang-format the output file")
    parser.add_argument("-o", help="Folder to write the tests to")
    args = parser.parse_args()

    with args.input as f:
        mapping = json.load(f)
    keys = list(mapping.keys())
    keys.sort()
    for k in keys:
        if ignore_function(k):
            continue
        overloads = mapping[k]
        kind, builtin = get_builtin_name(k)
        test_file = os.path.join(args.o, kind, builtin  + ".cl")
        if not os.path.exists(os.path.join(args.o, kind)):
            os.makedirs(os.path.join(args.o, kind))
        with open(test_file, "w") as out_fd:
            print("gen: " + test_file)
            out_fd.write("""
//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

// Autogenerated by gen-libclc-test.py

// RUN: %clang -emit-llvm -S -o - %s | FileCheck %s

#include <spirv/spirv_types.h>

// CHECK-NOT: declare {{.*}} @_Z
// CHECK-NOT: call {{[^ ]*}} bitcast
""")
            for proto, _ in overloads:
                if ignore_overload(proto):
                    continue
                nb_guards = emit_guards(out_fd, proto)
                ret = proto[0]
                proto = proto[1:]
                param = {
                    "RET" : ret,
                    "FN"  : k,
                    "PARAM" : ", ".join(["{} args_{}".format(ty, str(i)) for i, ty in enumerate(proto)]),
                    "RETURN" : "return" if ret != "void" else "",
                    "ARG" : ", ".join(["args_{}".format(str(i)) for i, _ in enumerate(proto)])
                }
                out_fd.write("""__attribute__((overloadable)) {RET} test_{FN}({PARAM}) {{
  {RETURN} {FN}({ARG});
}}

""".format(**param))
                close_guards(out_fd, nb_guards)
        if args.format:
            clang_format(args.format, test_file)