import re
import argparse


def get_platforms(file_name):
    # removing the new line characters
    with open(file_name) as f:
        lines = [line.rstrip() for line in f]

    platform_map = {}
    platform_lines = ""
    function_line = ""
    for line in lines:
        if "@platform" in line:
            platform_lines += line.replace("*", "").lstrip()
        if platform_lines == "":
            continue
        l = line.lstrip()
        if len(l) == 0:
            continue
        if l[0] == "*" or l[0] == "#" or (len(l) > 1 and l[0] == "/" and l[1] == "*"):
            if function_line != "":
                func = " ".join(function_line.split())
                platforms = re.findall(r"\{(.*?)\}", platform_lines)
                for p in platforms:
                    if p not in platform_map:
                        platform_map[p] = [func]
                    else:
                        platform_map[p].append(func)
                platform_lines = ""
            function_line = ""
            continue
        function_line += line + " "
    return platform_map


# string in list1 but not in list2
def diff(list1, list2):
    result = []
    for l1 in list1:
        if l1 not in list2:
            result.append(l1)
    return result


def common(list1, list2):
    result = []
    for l1 in list1:
        if l1 in list2:
            result.append(l1)
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="amdsmi platform classifier")
    parser.add_argument(
        "--platforms", default=False, action="store_true", help="list supported platforms"
    )
    parser.add_argument(
        "--list", default="", metavar="platform", help="List function in a platform"
    )
    parser.add_argument(
        "--diff",
        default=None,
        metavar="platform",
        nargs=2,
        help="Find the APIs in platform1 but not in platform2",
    )
    parser.add_argument(
        "--common",
        default=None,
        metavar="platform",
        nargs=2,
        help="Find the common APIs of two platforms",
    )

    file_name = "../include/amd_smi/amdsmi.h"
    args = parser.parse_args()
    platform_map = get_platforms(file_name)
    if args.platforms:
        for p in platform_map:
            print(p)
        exit(0)

    if args.list:
        if args.list not in platform_map:
            print("Unknown platform ", args.list)
            exit(1)
        for f in platform_map[args.list]:
            print(f)
        exit(0)

    if args.diff != None:
        if args.diff[0] not in platform_map or args.diff[1] not in platform_map:
            print("Unknown platforms ", args.diff)
            exit(1)
        platforms0 = platform_map[args.diff[0]]
        platforms1 = platform_map[args.diff[1]]
        result = diff(platforms0, platforms1)
        print("APIs in", args.diff[0], "but not in", args.diff[1])
        for f in result:
            print(f)
        exit(0)

    if args.common != None:
        if args.common[0] not in platform_map or args.common[1] not in platform_map:
            print("Unknown platforms ", args.common)
            exit(1)
        platforms0 = platform_map[args.common[0]]
        platforms1 = platform_map[args.common[1]]
        result = common(platforms0, platforms1)
        print("APIs in both ", args.common[0], "and", args.common[1])
        for f in result:
            print(f)
        exit(0)

    parser.print_help()