#! /usr/bin/env python3

"""
Compare two build trees to see if compiler or link flags have changed.
"""

import argparse, sys, subprocess, glob
from pathlib import Path

###############################################################################
def run_cmd_no_fail(cmd, from_dir=None):
###############################################################################
    proc = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                            cwd=from_dir, text=True)
    output, errput = proc.communicate()
    stat = proc.wait()

    assert stat == 0, f"CMD: {cmd} FAILED when run from {from_dir}\nERROR: {errput}"
    assert isinstance(output, str)

    return output.strip()

###############################################################################
def parse_command_line(args, description):
###############################################################################
    parser = argparse.ArgumentParser(
        usage="""\n{0} <filename>
OR
{0} --help

\033[1mEXAMPLES:\033[0m
    \033[1;32m# Compare case1 to case2 \033[0m
    > {0} $case1dir $case2dir

    \033[1;32m# same as ^, except with comparisons not sensitive to flag order. This flag is likely to be necessary for the flags.make to match \033[0m
    > {0} $case1dir $case2dir -u

    \033[1;32m# same as ^, except limit to atm component \033[0m
    > {0} $case1dir $case2dir -u -l atm.dir

    \033[1;32m# same as ^, except limit to link flags for all components \033[0m
    > {0} $case1dir $case2dir -u -l link.txt
""".format(Path(args[0])),
        description=description,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument("case1", help="The path to case1")

    parser.add_argument("case2", help="The path to case2")

    parser.add_argument("-u", "--unordered", action="store_true",
                        help="Make comparisons not sensitive to token order")

    parser.add_argument("-l", "--limit", dest="limits", action="append", default=[],
                        help="Limit compared files to files containing this substring. This option can be supplied multiple times")

    parser.add_argument("-b", "--build-log", action="store_true", help="Compare build log files instead of cmake files. This works best when e3sm is built with -j1")

    parser.add_argument("-v", "--verbose", action="store_true", help="Print diff details")

    args = parser.parse_args(args[1:])

    return args

###############################################################################
def find_files(builddir_cmake, filename, limits):
###############################################################################
    result = []
    all_files = list(sorted(glob.glob(f"{builddir_cmake}/**/{filename}", recursive=True)))

    if limits:
        for curr_file in all_files:
            matches_limits = True
            for limit in limits:
                if limit not in curr_file:
                    matches_limits = False
                    print(f"File {curr_file} does not match limit {limit}, skipping")
                    break

            if matches_limits:
                result.append(curr_file)

        return result

    else:
        return all_files

###############################################################################
def compare_tokens(line1, line2, unordered):
###############################################################################
    tokens1 = line1.split()
    tokens2 = line2.split()

    if unordered:
        tokens1.sort()
        tokens2.sort()

    result = True
    message = ""
    if tokens1 != tokens2:
        result = False

        tokens1_set = set(tokens1)
        tokens2_set = set(tokens2)
        one_not_two = tokens1_set - tokens2_set
        two_not_one = tokens2_set - tokens1_set
        #message += f"    For line {line1}:"
        if one_not_two:
            message += "      Case1 had unique flags:\n"
            for item in one_not_two:
                message += f"        {item}\n"

        if two_not_one:
            message += "      Case2 had unique flags:\n"
            for item in two_not_one:
                message += f"        {item}\n"

    return result, message

###############################################################################
def compare_contents(case1, case2, file1, file2, contents1, contents2, unordered, verbose):
###############################################################################
    print("###############################################################################")
    print(f"COMPARING FILES {file1} AND {file2}")

    normalized_contents2 = contents2.replace(case2, case1)

    lines1 = contents1.splitlines()
    lines2 = normalized_contents2.splitlines()

    assert len(lines1) == len(lines2), f"{file1} and {file2} are not even the same length!"

    files_match = True
    for line1, line2 in zip(lines1, lines2):
        result, message = compare_tokens(line1, line2, unordered)
        if verbose:
            print(message)
        files_match &= result

    if files_match:
        print("\n  FILES MATCHED")
    else:
        print("\n  FILES DID NOT MATCH")

    return files_match

###############################################################################
def compare_file_lists(case1, case2, files1, files2, unordered, verbose):
###############################################################################
    result = True
    for file1, file2 in zip(files1, files2):
        assert file1 == file2, f"File orders did not match, {file1} != {file2}"

        file1c = Path(file1).open("r", encoding="utf8").read()
        file2c = Path(file2).open("r", encoding="utf8").read()

        result &= compare_contents(case1, case2, file1, file2, file1c, file2c, unordered, verbose)

    return result

###############################################################################
def get_log_from_blddir(builddir):
###############################################################################
    logs = glob.glob(f"{builddir}/e3sm.bldlog.*")
    assert len(logs) == 1, f"Expected to match exactly one log file, got {logs}"

    logf = Path(logs[0])
    assert logf.suffix != ".gz", f"Expected {logf} to be uncompressed"

    return logf

###############################################################################
def parse_log(logf, normalize=None):
###############################################################################
    log_content = logf.open("r", encoding="utf8").read()

    result = {}
    for line in log_content.splitlines():
        if "e3sm_compile_wrap.py" in line:
            if normalize is not None:
                for repl, repl_with in normalize:
                    line = line.replace(repl, repl_with)

            args = line.split()
            arglen = len(args)
            target = None
            for idx, arg in enumerate(args):
                if arg == "-o" and idx + 1 < arglen:
                    target = args[idx + 1]
                    break

                if arg.startswith("lib") and arg.endswith(".a"):
                    if args[idx-1].endswith("ranlib"):
                        target = "skipme"
                    else:
                        target = arg

            if target is None:
                target = line
                print(f"WARNING: Failed to parse line {line}")

            if target != "skipme":
                assert target not in result, f"{target} appears twice in log {logf}?\nFirst encountered: {result[target]}\n\nAnd again: {line}"
                result[target] = line

    return result

###############################################################################
def compare_logs(builddir1, builddir2, casename1, casename2, srcroot1, srcroot2, unordered, verbose):
###############################################################################
    result = True

    log1 = get_log_from_blddir(builddir1)
    log2 = get_log_from_blddir(builddir2)

    log1c = parse_log(log1)
    log2c = parse_log(log2, normalize=((casename2, casename1), (srcroot2, srcroot1)))

    message_dict = {}

    for target, line1 in log1c.items():
        if target in log2c:
            line2 = log2c[target]
            print(f"Comparing target {target}")

            curr_result, message = compare_tokens(line1, line2, unordered)
            if verbose and not curr_result:
                if message in message_dict:
                    print(f"  Same issue as {message_dict[message]}")
                else:
                    message_dict[message] = target
                    print(message)

            result &= curr_result
            if curr_result:
                print(f"  Target has identical build")
            else:
                print(f"  Target did not have identical build")

        else:
            print(f"WARNING: target {target} is missing from case2")
            result = False

    return result

###############################################################################
def compare_flags(case1, case2, unordered, limits, build_log, verbose):
###############################################################################
    result = True

    case1 = Path(case1)
    case2 = Path(case2)

    assert case1.is_dir(), f"{case1} is not an existing directory"
    assert case2.is_dir(), f"{case2} is not an existing directory"

    builddir1 = Path(run_cmd_no_fail("./xmlquery EXEROOT --value", from_dir=case1))
    builddir2 = Path(run_cmd_no_fail("./xmlquery EXEROOT --value", from_dir=case2))

    casename1 = run_cmd_no_fail("./xmlquery CASE --value", from_dir=case1)
    casename2 = run_cmd_no_fail("./xmlquery CASE --value", from_dir=case2)

    srcroot1 = run_cmd_no_fail("./xmlquery SRCROOT --value", from_dir=case1)
    srcroot2 = run_cmd_no_fail("./xmlquery SRCROOT --value", from_dir=case2)

    builddir1_cmake = builddir1 / "cmake-bld"
    builddir2_cmake = builddir2 / "cmake-bld"

    assert builddir1_cmake.is_dir(), \
        f"{builddir1_cmake} is not an existing directory, you need to run case.build in {case1}"
    assert builddir2_cmake.is_dir(), \
        f"{builddir2_cmake} is not an existing directory, you need to run case.build in {case2}"

    if build_log:
        result = compare_logs(builddir1, builddir2, casename1, casename2, srcroot1, srcroot2, unordered, verbose)

    else:
        flag_files1 = find_files(builddir1_cmake, "flags.make", limits)
        flag_files2 = find_files(builddir2_cmake, "flags.make", limits)

        link_files1 = find_files(builddir1_cmake, "link.txt", limits)
        link_files2 = find_files(builddir2_cmake, "link.txt", limits)

        print()

        result &= compare_file_lists(casename1, casename2, flag_files1, flag_files2, unordered, verbose)

        result &= compare_file_lists(casename1, casename2, link_files1, link_files2, unordered, verbose)

    if result:
        print("\nALL FILES MATCHED")
    else:
        print("\nALL FILES DID NOT MATCH")

    return result

###############################################################################
def _main_func(description):
###############################################################################
    success = compare_flags(**vars(parse_command_line(sys.argv, description)))

    sys.exit(0 if success else 1)

###############################################################################

if (__name__ == "__main__"):
    _main_func(__doc__)
