#!/usr/bin/env python3
"""
Wrapper compiler based on GPUFORT's MACA C++ code generation capabilities.
"""
import os
import re
import sys
import argparse
import textwrap
from typing import List
from pathlib import Path
sys.path.append(str(Path(__file__).resolve().parents[2]))

from mcgpufort import gpufort
from mcgpufort import utils
from mcgpufort import options as opts
script_dir = os.path.dirname(__file__)
MACA_PATH = os.getenv("MACA_PATH")

def populate_cl_arg_parser(parser: argparse.ArgumentParser):
    # General options
    parser.add_argument("inputs",
                        help="The input files",
                        type=str,
                        nargs="+",
                        default=[])
    parser.add_argument("-c",
                        dest="only_compile",
                        action="store_true",
                        help="Only compile; do not link")
    parser.add_argument("-o",
                        dest="output",
                        default=None,
                        type=str,
                        help="Path/name for the result.")
    parser.add_argument(
        "--gpufort-fc",
        dest="fortran_compiler",
        default="gfortran",
        type=str,
        help="Compiler for Fortran code, defaults to 'gfortran'. \
              Use '--gpufort-fcflags <flags>' to set compilation flags and \
              '--gpufort-ldflags <flags>' to set linker flags. Flags to GPUFORT \
              related libraries are added by default if not specified otherwise. \
              Host and MACA C++ compiler and their respective options must be \
              specified after all other GPUFORT options."
    )
    parser.add_argument(
        "--gpufort-cc",
        dest="cpp_compiler",
        default="mxcc",
        type=str,
        help="Compiler for MACA C++ code, defaults to 'mxcc'. \
              Use '--gpufort-cflags <flags>' to set compilation flags and \
              '--gpufort-ldflags <flags>' to set linker flags. Flags to GPUFORT \
              related libraries are added by default if not specified otherwise. \
              Host and MACA C++ compiler and their respective options must be \
              specified after all other GPUFORT options."
    )
    parser.add_argument(
        "--gpufort-ldflags",
        dest="gpufort_ldflags",
        default="",
        type=str,
        help="Linker flags"
    )
    parser.add_argument(
        "-save-temps",
        dest="save_temps",
        action="store_true",
        help="Save temporary output files, i.e. intermediate results produced \
              by the translation of the original code to standard Fortran plus \
              MACA C++."
    )
    parser.add_argument(
        "-only-codegen",
        dest="only_codegen",
        action="store_true",
        help="Only generate code. Do not compile."
    )
    parser.add_argument(
        "-no-codegen",
        dest="no_codegen",
        action="store_true",
        help="Do not generate code. This assumes that files has been done before."
    )
    parser.add_argument(
        "-no-fallback",
        dest="no_fallback",
        action="store_true",
        help="Do not fallback to compile original fortran file."
    )
    parser.set_defaults(only_compile=False,
                        only_codegen=False,
                        no_codegen=False,
                        no_fallback=False)

def check_inputs(paths: list):
    """
    :return: Typical object file name/path/ that compilers derive from a
        Fortran/C++ sources file.
    :raise ValueError: If no single file with Fortran file extension could
        be identified and other inputs do not have an '.so', '.a', or '.o'
        ending.
    """
    fortran_files = list()
    objects       = list()
    others        = list()
    for path in paths:
        abspath = os.path.abspath(path)
        if re.match(r".+\.(f[0-9]*)$", path, re.IGNORECASE):
            fortran_files.append(abspath)
        elif re.match(r".+\.(so|a|o)$", path, re.IGNORECASE):
            objects.append(abspath)
        else:
            others.append(path)
    if len(fortran_files) > 1:
        raise ValueError(
            "specify only a single Fortran file as input; have: " +
            ",".join(fortran_files)
        )
    if len(others):
        raise ValueError(
            "The following files are neither Fortran nor object files:" +
            ",".join(others)
        )
    return fortran_files[0], objects

def get_object_path(source_path: str):
    """
    :return: Typical object file name/path that compilers derive from a
        Fortran/C++ source file.
    """
    result = re.sub(r"\.(f[0-9]*|F[0-9]*)$", ".o", source_path, re.IGNORECASE)
    result = re.sub(r"\.(c|cc|cxx|cpp)$", ".cpp.o", result, re.IGNORECASE)
    return result

def handle_subprocess_output(triple: tuple, exit: bool=True):
    status, output, errors = triple
    sys.stdout.write(output)
    sys.stderr.write(errors)
    if status and exit:
        sys.exit(status)
    return status

def hardcore_pure_gpufort_opts():
    """Fixes certain opts, makes them unaffected by config."""
    pass

def remove_file(path):
    """Tries to remove file. Throws no error if it doesn't exists."""
    try:
        os.remove(path)
    except OSError:
        pass

def recored_fallback(file_name: str):
    file_path = os.environ.get("ACC_RECORED_FALLBACK")
    if file_path:
        print(f'fallback file:{file_name}')
        file_path = os.path.join(file_path, "fallback.txt")
        with open(file_path, "a")as file:
            file.write("\n")
            file.write(file_name)

def fallback(infile_path: str,
             include_dirs: List[str],
             compiler_options: List[str],
             args: argparse.Namespace):
    ld_flags = list()
    ld_flags += gpufort.get_basic_ldflags(True)
    ld_flags += gpufort.get_mxcc_ldflangs()
    include_dirs.append(gpufort.get_gpufort_include())
    outfile_path = args.output
    if outfile_path is None:
        if args.only_compile:
            outfile_path = get_object_path(infile_path)
        else:
            outfile_path = os.path.join(args.working_dir, "a.out")
    cmd = list()
    cmd.append(args.fortran_compiler)
    if args.only_compile:
        cmd.append("-c")
    cmd.extend(compiler_options)
    cmd.extend(map(lambda inc: f"-I{inc}", include_dirs))
    cmd.extend(["-o", outfile_path])
    cmd.append(infile_path)
    cmd.extend(map(lambda inc: f"{inc}", ld_flags))
    print("mxfortran start fallback")
    status, output, errors = utils.subprocess.run_subprocess(tuple(cmd))
    recored_fallback(infile_path)
    sys.stdout.write(output)
    sys.stderr.write(errors)
    sys.exit(status)

if __name__ == "__main__":
    # read config and command line arguments
    include_dirs, defines, compiler_options = gpufort.parse_raw_command_line_arguments()
    # parse command line arguments
    parser = argparse.ArgumentParser(description=textwrap.dedent("""
Wrapper compiler for OpenACC Fortran programs based on GPUFORT's MACA
C++ code generator.

NOTE: This tool is based on GPUFORT tool. Options inherited from the
latter have a '--' prefix while this tool's option have a '-' prefix."""))
    populate_cl_arg_parser(parser)
    gpufort.populate_cl_arg_parser(parser, False)
    args, unknown_args = gpufort.parse_cl_args(parser, False)

    if args.only_codegen and args.no_codegen:
        utils.logging.log_error(
            opts.log_prefix, "__main__",
            "'-only-codegen' and '-no-codegn' are mutually exclusive"
        )
        sys.exit(2)

    try:
        infile_path, objects_files = check_inputs(args.inputs)
    except ValueError as e:
        utils.logging.log_error(opts.log_prefix, "__main__", str(e))
        sys.exit(2)
    gpufort.init_logging(args.inputs[0])

    gpufort.set_include_dirs(args.working_dir, include_dirs)
    gpufort.map_args_to_opts(args, opts.include_dirs, defines, False)
    hardcore_pure_gpufort_opts()

    modified_fortran_file_path = "{}/{}{}".format(args.working_dir,
                                                  os.path.basename(infile_path),
                                                  opts.fortran_file_ext)
    cpp_file_path = "{}/{}{}".format(args.working_dir,
                                     os.path.basename(infile_path),
                                     gpufort.fort2maca.opts.cpp_file_ext)
    fortran_module_file_path = "{}/{}{}".format(
        args.working_dir,
        os.path.basename(infile_path),
        gpufort.fort2maca.opts.fortran_module_file_ext
    )
    is_generate_kernel_files = False
    if args.only_codegen or not args.no_codegen:
        is_generate_kernel_files = gpufort.run_checked(
            infile_path,
            modified_fortran_file_path,
            cpp_file_path,
            defines,
            raise_exception=args.no_fallback
        )
        if not (is_generate_kernel_files or args.only_codegen):
            gpufort.shutdown_logging()
            fallback(infile_path, include_dirs, compiler_options, args)

    # output file name
    if not args.only_codegen:
        outfile_path = args.output
        if outfile_path is None:
            if args.only_compile:
                outfile_path = get_object_path(infile_path)
            else:
                outfile_path = os.path.join(args.working_dir, "a.out")

        fc_path = ""
        cc_path = ""
        # flags
        fortran_cflags = list()
        cpp_cflags     = list()
        ldflags        = list()
        if args.use_default_flags:
            cpp_cflags += gpufort.get_basic_cflags()
            fortran_cflags += list(cpp_cflags)
            ldflags += gpufort.get_basic_ldflags(args.openacc)
            if os.path.basename(args.fortran_compiler) == "gfortran":
                fortran_cflags += gpufort.get_gfortran_cflags()
                ldflags += gpufort.get_gfortran_ldflangs()
                fc_path += ':/usr/bin'
            if os.path.basename(args.fortran_compiler) == "flang-new":
                ldflags += ['-L' + os.path.join(script_dir, "../../flang17/lib")]
                ldflags += gpufort.get_flang_ldflangs()
                fc_path += f':{script_dir}/../../flang17/bin:{script_dir}/../flang17/bin'
            if os.path.basename(args.cpp_compiler) == "mxcc":
                cpp_cflags += gpufort.get_mxcc_cflags()
                ldflags += gpufort.get_mxcc_ldflangs()
                cc_path += f':{MACA_PATH}/mxgpu_llvm/bin'
        ldflags.insert(0,args.gpufort_ldflags)
        # cpp compilation
        cpp_object_path = get_object_path(cpp_file_path)
        cpp_cmd = tuple([
            args.cpp_compiler,
            "-c",
            *cpp_cflags,
            "-o", cpp_object_path,
            cpp_file_path
        ])
        status = handle_subprocess_output(
            utils.subprocess.run_subprocess(
                cpp_cmd,
                args.verbose,
                {"PATH": os.getenv("PATH")+f'{cc_path}'}
            ),
            exit=args.no_fallback
        )
        if status and not args.no_fallback:
            gpufort.shutdown_logging()
            remove_file(cpp_file_path)
            remove_file(fortran_module_file_path)
            remove_file(modified_fortran_file_path)
            fallback(infile_path, include_dirs, compiler_options, args)

        # fortran compilation
        modified_object_path = get_object_path(modified_fortran_file_path)
        fortran_cmd = tuple([
            args.fortran_compiler,
            "-c",
            *fortran_cflags,
            *compiler_options,
            "-o", modified_object_path,
            modified_fortran_file_path
        ])
        status = handle_subprocess_output(
            utils.subprocess.run_subprocess(
                fortran_cmd,
                args.verbose,
                {"PATH": os.getenv("PATH")+f'{fc_path}'}
            ),
            exit=args.no_fallback
        )
        if status and not args.no_fallback:
            gpufort.shutdown_logging()
            remove_file(cpp_file_path)
            remove_file(cpp_object_path)
            remove_file(fortran_module_file_path)
            remove_file(modified_fortran_file_path)
            fallback(infile_path, include_dirs, compiler_options, args)
        
        # emit a signle object combining the Fortran and C++ objects
        if args.only_compile:
            # merge object files; produces: outfile_path
            merge_cmd = tuple([
                "/usr/bin/ld",
                "-r",
                "-o", outfile_path,
                modified_object_path, cpp_object_path
            ])
            handle_subprocess_output(
                utils.subprocess.run_subprocess(merge_cmd, args.verbose))
        # emit an executable
        else:
            binary_cmd = tuple([
                args.cpp_compiler,
                *ldflags,
                "-o", outfile_path,
                modified_object_path, cpp_object_path
            ])
            handle_subprocess_output(
                utils.subprocess.run_subprocess(binary_cmd, args.verbose, {"PATH": os.getenv("PATH")+f'{cc_path}'}))

        if not args.save_temps:
            remove_file(cpp_file_path)
            remove_file(modified_fortran_file_path)
            remove_file(cpp_object_path)
            remove_file(modified_object_path)
            remove_file(fortran_module_file_path)

    gpufort.shutdown_logging()
