#!/usr/bin/env python3

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


config = A.config()
prefix = Path(config["prefix"])
CXX = Path(config["CXX"])
if not CXX.exists():
    # Example <>/lib/python3.10/site-packages/arbor
    altern = "c++"
    from shutil import which
    if which(altern) is None:
        print(
            f"Error: Neither prefix '{CXX}' nor fallback '{altern}' exist. Please provide a path to a C++ compiler with --cxx",
            file=sys.stderr,
        )
        exit(-1)
    print(
        f"Warning: prefix '{CXX}' does not exist, falling back to '{altern}'.",
        file=sys.stderr,
    )
    CXX = altern


def parse_arguments():
    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=argparse.ArgumentDefaultsHelpFormatter,
    )

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

    parser.add_argument(
        "--raw",
        metavar="raw",
        nargs="+",
        default=[],
        type=str,
        help="""Raw C++ mechanisms; per <name> the files <name>.hpp,
<name>_cpu.cpp (if CPU enabled), <name>_gpu.cpp and <name>_gpu.cu (if GPU
enabled) must be present in the target directory.""",
    )

    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("--cpu", default=True, help="Enable CPU support.")

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

    parser.add_argument(
        "--gpu",
        metavar="gpu",
        default=config["gpu"] if config["gpu"] else "none",
        choices=["none", "cuda", "hip", "cuda-clang"],
        help="Enable GPU support",
    )

    parser.add_argument(
        "--cxx",
        metavar="cxx",
        default=CXX,
        help="Use this C++ compiler.",
    )

    parser.add_argument(
        "--prefix",
        metavar="prefix",
        default=prefix,
        help="Arbor's install prefix.",
    )

    parser.add_argument(
        "--bin",
        metavar="bin",
        default=config["binary_path"],
        help="Look here for Arbor utils like modcc; relative to prefix.",
    )

    parser.add_argument(
        "--lib",
        metavar="lib",
        default=config["lib_path"],
        help="Look here for Arbor's CMake config; relative to prefix.",
    )

    parser.add_argument(
        "--data",
        metavar="data",
        default=config["data_path"],
        help="Look here for Arbor supplementals like generate_catalogue; relative to prefix.",
    )

    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
if verbose:
    print("=" * 80)
    print(f"{os.path.basename(__file__)} called with the following arguments:")
    for k, v in args.items():
        print(k, v)
    print("=" * 80)
debug = args["debug"]
raw = args["raw"]
gpu = args["gpu"]
cpu = args["cpu"]

if gpu == "cuda":
    gpu_support = f"""
include(FindCUDAToolkit)
add_compile_definitions(ARB_CUDA)
add_compile_definitions(ARB_HAVE_GPU)
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
    # Pascal, Volta, Ampere, Hopper
    set(CMAKE_CUDA_ARCHITECTURES 60 70 80 90)
endif()
find_package(CUDAToolkit)
enable_language(CUDA)
set(CMAKE_CUDA_STANDARD 20)
set(CMAKE_CUDA_HOST_COMPILER {args["cxx"]})
"""
elif gpu == "cuda-clang":
    print("CUDA-Clang support is currently considered experimental only.")
    gpu_support = """
add_compile_definitions(ARB_CUDA)
add_compile_definitions(ARB_HAVE_GPU)
find_package(CUDAToolkit)
enable_language(CUDA)
set(CMAKE_CUDA_STANDARD 20)
add_compile_options(-xcuda --cuda-gpu-arch=sm_60 --cuda-gpu-arch=sm_70 --cuda-gpu-arch=sm_80 --cuda-path=${CUDA_TOOLKIT_ROOT_DIR}))
"""
elif gpu == "hip":
    print("HIP support is currently considered experimental only.")
    gpu_support = """
add_compile_definitions(ARB_HIP)
add_compile_definitions(ARB_HAVE_GPU)
add_compile_options(-xhip --amdgpu-target=gfx906 --amdgpu-target=gfx900)
"""
elif gpu == "none":
    gpu_support = """
# GPU: Disabled
"""
else:
    print(f"Internal Error: Unknown GPU type: {gpu}", file=sys.stderr)
    exit(-1)

paths = [
    Path(args["prefix"]),
    Path(A.__path__[0]).parent.parent.parent.parent,
    Path(A.__path__[0]),
]

for path in paths:
    reqs = [path / args["bin"] / "modcc",
            path / args["data"] / "arbor" / "BuildModules.cmake",
            path / args["lib"] / "cmake" / "arbor" / "arbor-config.cmake",]

    if all(req.exists() for req in reqs):
        prefix = path
        break
else:
    raise FileNotFoundError("Could not find required tools. Please check your installation.")

bindir, datdir, pakdir = prefix / args['bin'], path / args["data"] / "arbor", path / args["lib"] / "cmake"

cmake = f"""
cmake_minimum_required(VERSION 3.9)
project({name}-cat LANGUAGES CXX)
set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)

list(APPEND CMAKE_PREFIX_PATH {pakdir}/arbor)
list(APPEND CMAKE_PREFIX_PATH {pakdir}/units)
find_package(arbor REQUIRED)
{gpu_support}
set(CMAKE_BUILD_TYPE release)
set(CMAKE_CXX_COMPILER  {args["cxx"]})
set(CMAKE_CXX_FLAGS     ${{ARB_CXX_FLAGS}})

include(BuildModules.cmake)

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

make_catalogue_standalone(
  NAME {name}
  SOURCES "${{CMAKE_CURRENT_SOURCE_DIR}}/mod"
  MOD {' '.join(mods)}
  CXX {' '.join(raw)}
  CXX_FLAGS_TARGET ${{ARB_CXX_FLAGS_TARGET}}
  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, **_):
            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, **_):
            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"{datdir}/BuildModules.cmake", tmp)

    out = tmp / "build" / "generated" / name
    os.makedirs(out, exist_ok=True)
    sfx = [".hpp"]
    if cpu:
        sfx += ["_cpu.cpp"]
    if gpu and gpu != 'none':
        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:
        make_cmd += " VERBOSE=1"
    try:
        sp.run(cmake_cmd, shell=True, check=True, capture_output = True, text = True)
        sp.run(make_cmd, shell=True, check=True, capture_output = True, text = True)
        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)
            sys.stderr.write(tb.format_exc() + " Error:\n\n")
            sys.stderr.write(e.stderr)
        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")
