#!/usr/bin/env python3

import subprocess as sp
import sys
from tempfile import mkdtemp
import os
from pathlib import Path
import shutil
import argparse
import re


def parse_arguments():
    def append_slash(s):
        return s + "/" if s and not s.endswith("/") else s

    class ConciseHelpFormatter(argparse.HelpFormatter):
        def __init__(self, **kwargs):
            super(ConciseHelpFormatter, self).__init__(max_help_position=20, **kwargs)

        def _format_action_invocation(self, action):
            if not action.option_strings:
                return super(ConciseHelpFormatter, self)._format_action_invocation(
                    action
                )
            else:
                optstr = ", ".join(action.option_strings)
                if action.nargs == 0:
                    return optstr
                else:
                    return optstr + " " + self._format_args(action, action.dest.upper())

    parser = argparse.ArgumentParser(
        description="Generate dynamic catalogue and build it into a shared object.",
        usage="%(prog)s catalogue_name mod_source_dir",
        add_help=False,
        formatter_class=ConciseHelpFormatter,
    )

    parser.add_argument("name", metavar="name", type=str, help="Catalogue name.")

    parser.add_argument(
        "--raw",
        metavar="raw",
        nargs="+",
        default=[],
        type=str,
        help="""Advanced: Raw mechanisms as C++ files. Per <name> the
files <name>.hpp, <name>_cpu.cpp must be present
in the target directory and with GPU support
also <name>_gpu.cpp and <name>_gpu.cu (if not given -C).""",
    )

    parser.add_argument(
        "modpfx",
        metavar="modpfx",
        type=str,
        help="Directory name where *.mod files live.",
    )

    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose.")

    parser.add_argument("-q", "--quiet", action="store_true", help="Less output.")

    parser.add_argument(
        "-g",
        "--gpu",
        metavar="gpu",
        help="Enable GPU support, valid options: cuda|hip|cuda-clang.",
    )

    parser.add_argument(
        "-C", "--no-cpu", action="store_true", help="Disable CPU support."
    )

    parser.add_argument(
        "-d",
        "--debug",
        nargs="?",
        metavar="path",
        const=True,
        default=False,
        help="Don't clean up the generated temp cpp code."
        + " Can be a target path for the generated code.",
    )

    parser.add_argument(
        "-h", "--help", action="help", help="Display this help and exit."
    )

    return vars(parser.parse_args())


args = parse_arguments()
pwd = Path.cwd()
name = re.sub(r"_+", r"_", re.sub(r"[^a-zA-Z0-9_]", r"_", args["name"]))

mod_dir = pwd / Path(args["modpfx"])
mods = [f[:-4] for f in os.listdir(mod_dir) if f.endswith(".mod")]
quiet = args["quiet"]
verbose = args["verbose"] and not quiet
debug = args["debug"]
raw = args["raw"]
gpu = args["gpu"]
cpu = not args["no_cpu"]

if gpu:
    if gpu == "cuda":
        gpu_support = """
add_compile_definitions(ARB_CUDA)
add_compile_definitions(ARB_HAVE_GPU)

enable_language(CUDA)
set(CMAKE_CUDA_HOST_COMPILER /Applications/Xcode_13.2.1.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/bin/c++)
set(CMAKE_CUDA_ARCHITECTURES )
"""
    else:
        print(
            f"Unsupported GPU target: {gpu}. If you need support for HIP or Clang-CUDA, please check here: https://github.com/arbor-sim/arbor/issues/1783"
        )
        exit(-1)
else:
    gpu_support = """
# GPU: Disabled
"""

this_path = Path(__file__).parent
data_path = (this_path / "../share/arbor").resolve()
pack_path = (this_path / "../lib/cmake/arbor").resolve()
exec_path = this_path.resolve()

for path in [
    exec_path / "modcc",
    data_path / "generate_catalogue",
    data_path / "BuildModules.cmake",
    pack_path / "arbor-config.cmake",
]:
    if not path.exists():
        print(f"Could not find required tool: {path}. Please check your installation.")
        exit(-1)

cmake = f"""
cmake_minimum_required(VERSION 3.9)
project({name}-cat LANGUAGES CXX)

set(arbor_DIR {pack_path})
find_package(arbor REQUIRED)
{gpu_support}
set(CMAKE_BUILD_TYPE release)
set(CMAKE_CXX_COMPILER  ${{ARB_CXX}})
set(CMAKE_CXX_FLAGS     ${{ARB_CXX_FLAGS}})

include(BuildModules.cmake)

set(ARB_WITH_EXTERNAL_MODCC true)
find_program(modcc NAMES modcc PATHS {exec_path})

make_catalogue(
  NAME {name}
  SOURCES "${{CMAKE_CURRENT_SOURCE_DIR}}/mod"
  OUTPUT "CAT_{name.upper()}_SOURCES"
  MOD {' '.join(mods)}
  CXX {' '.join(raw)}
  PREFIX {data_path}
  CXX_FLAGS_TARGET ${{ARB_CXX_FLAGS_TARGET}}
  STANDALONE ON
  VERBOSE {"ON" if verbose else "OFF"})
"""

if not quiet:
    print(f"Building catalogue '{name}' from mechanisms in {mod_dir}")
    if debug:
        print("Debug mode enabled.")
    if mods:
        print(" * NMODL")
        for m in mods:
            print("   *", m)
    if raw:
        print(" * Raw")
        for m in raw:
            print("   *", m)

if debug:
    # Overwrite the local reference to `TemporaryDirectory` with a context
    # manager that doesn't clean up the build folder so that the generated cpp
    # code can be debugged
    class TemporaryDirectory:
        def __enter__(*args, **kwargs):
            if isinstance(debug, str):
                path = os.path.abspath(debug)
                try:
                    os.makedirs(path, exist_ok=False)
                except FileExistsError:
                    sys.stderr.write(
                        f"Error: Debug destination '{path}' already exists.\n"
                    )
                    sys.stderr.flush()
                    exit(1)
            else:
                path = mkdtemp()
            print(f"Building debug code in '{path}'.")
            return path

        def __exit__(*args, **kwargs):
            pass

else:
    from tempfile import TemporaryDirectory

with TemporaryDirectory() as tmp:
    tmp = Path(tmp)
    shutil.copytree(mod_dir, tmp / "mod")
    os.mkdir(tmp / "build")
    os.chdir(tmp / "build")
    with open(tmp / "CMakeLists.txt", "w") as fd:
        fd.write(cmake)
    shutil.copy2(f"{data_path}/BuildModules.cmake", tmp)
    shutil.copy2(f"{data_path}/generate_catalogue", tmp)

    out = tmp / "build" / "generated" / name
    os.makedirs(out, exist_ok=True)
    sfx = [".hpp"]
    if cpu:
        sfx += ["_cpu.cpp"]
    if gpu:
        sfx += ["_gpu.cpp", "_gpu.cu"]
    for e in raw:
        for s in sfx:
            fn = mod_dir / (e + s)
            if not fn.exists():
                print(
                    f"Could not find required file: {fn}. Please check your C++ mechanisms."
                )
                exit(-1)
            else:
                shutil.copy2(fn, out / (e + s))

    cmake_cmd = "cmake .."
    make_cmd = "make"
    if verbose:
        out, err = (None, None)
        make_cmd += " VERBOSE=1"
    else:
        out, err = (sp.PIPE, sp.PIPE)
    try:
        sp.run(cmake_cmd, shell=True, check=True, stdout=out, stderr=err)
        sp.run(make_cmd, shell=True, check=True, stdout=out, stderr=err)
        shutil.copy2(f"{name}-catalogue.so", pwd)
    except sp.CalledProcessError as e:
        import traceback as tb

        if not verbose:
            # Not in verbose mode, so we have captured the
            # `stdout` and `stderr` and can print it to the user.
            sys.stdout.write("Build log:\n")
            sys.stdout.write(e.stdout.decode())
            sys.stderr.write(tb.format_exc() + " Error:\n\n")
            sys.stderr.write(e.stderr.decode())
        else:
            # In verbose mode the outputs weren't captured and
            # have been streamed to `stdout` and `stderr` already.
            sys.stderr.write(
                "Catalogue building error occurred."
                + " Check stdout log for underlying error,"
                + " or omit verbose flag to capture it."
            )
        sys.stdout.flush()
        sys.stderr.flush()
        exit(e.returncode)

    if not quiet:
        print(f"Catalogue has been built and copied to {pwd}/{name}-catalogue.so")
