#! /usr/bin/env python

import os
import getpass
import pathlib
import subprocess

import waves

# Comments used in tutorial code snippets: marker-1

# Accept command line options with fall back default values
AddOption(
    "--build-dir",
    dest="variant_dir_base",
    default="build",
    nargs=1,
    type="string",
    action="store",
    metavar="DIR",
    help="SCons build (variant) root directory. Relative or absolute path. (default: '%default')",
)
AddOption(
    "--unconditional-build",
    dest="unconditional_build",
    default=False,
    action="store_true",
    help="Boolean flag to force building of conditionally ignored targets. (default: '%default')",
)
AddOption(
    "--print-build-failures",
    dest="print_build_failures",
    default=False,
    action="store_true",
    help="Print task *.stdout target file(s) on build failures. (default: '%default')",
)
# Python optparse appends to the default list instead of overriding. Must implement default/override ourselves.
default_cubit_commands = [
    "/apps/Cubit-16.12/cubit",
    "/usr/projects/ea/Cubit/Cubit-16.12/cubit",
    "cubit",
]
AddOption(
    "--cubit-command",
    dest="cubit_command",
    nargs=1,
    type="string",
    action="append",
    metavar="COMMAND",
    help=f"Override for the Cubit command. Repeat to specify more than one (default: {default_cubit_commands})",
)

# Comments used in tutorial code snippets: marker-2

# Inherit user's full environment and set project variables
env = waves.scons_extensions.WAVESEnvironment(
    ENV=os.environ.copy(),
    variant_dir_base=pathlib.Path(GetOption("variant_dir_base")),
    unconditional_build=GetOption("unconditional_build"),
    print_build_failures=GetOption("print_build_failures"),
    cubit_commands=GetOption("cubit_command"),
)

# Conditionally print failed task *.stdout files
env.PrintBuildFailures(print_stdout=env["print_build_failures"])

# Always copy (no sym-links) when duplicating
env.SetOption("duplicate", "copy")

# Comments used in tutorial code snippets: marker-3

# Find required programs for conditional target ignoring and absolute path for use in target actions
env["CUBIT_PROGRAM"] = env.AddCubit(
    env["cubit_commands"] if env["cubit_commands"] is not None else default_cubit_commands
)
env["SBATCH_PROGRAM"] = env.AddProgram(["sbatch"])
env["user"] = getpass.getuser()

# Try a preference ordered list of project server-specific environment modulefiles
quinoa_hpc_modulefile = pathlib.Path("/usr/projects/ea/Quinoa/modules/modulefiles/quinoa")
quinoa_modulefile_options = [
    pathlib.Path("/projects/aea_compute/modulefiles/aea-quinoa"),
    quinoa_hpc_modulefile,
]
quinoa_modulefile = next((path for path in quinoa_modulefile_options if path.exists()), None)

# Comments used in tutorial code snippets: marker-4

# Set project internal variables and variable substitution dictionaries
project_name = "WAVES-TUTORIAL"
version = "0.1.0"
project_dir = pathlib.Path(Dir(".").abspath)
project_variables = {
    "project_name": project_name,
    "project_dir": project_dir,
    "version": version,
    "remote_directory": pathlib.Path(f"/users/{env['user']}/WAVES-TUTORIAL/tutorial_quinoa"),
}
for key, value in project_variables.items():
    env[key] = value

# Comments used in tutorial code snippets: marker-5

# Add builders and pseudo-builders
env.Append(
    BUILDERS={
        "SSHQuinoaSolver": waves.scons_extensions.ssh_builder_actions(
            waves.scons_extensions.sbatch_quinoa_builder_factory(
                environment=f"module use {quinoa_hpc_modulefile.parent} && module load {quinoa_hpc_modulefile.name} &&"
            ),
            remote_server="ro-rfe.lanl.gov",
        )
    }
)

# Source the Quinoa environment and store it separately from the Conda/Cubit environment to avoid PATH interference
if quinoa_modulefile:
    envQuinoa = waves.scons_extensions.shell_environment(
        "module unload aea-nightly aea-weekly aea-quarterly aea-release && "
        f"module use {quinoa_modulefile.parent} && "
        f"module load {quinoa_modulefile.name}"
    )
else:
    envQuinoa = Environment()
envQuinoa.AddMethod(waves.scons_extensions.add_program, "AddProgram")
envQuinoa["inciter"] = envQuinoa.AddProgram(["inciter"])
envQuinoa["charmrun"] = envQuinoa.AddProgram(["charmrun"])

# Add WAVES builders. If sbatch is found, use it.
# Should allow the same scons alias to run directly on sstbigbird, but submit as an sbatch job on HPC
quinoa_builder_factory = waves.scons_extensions.quinoa_builder_factory
if env["SBATCH_PROGRAM"] is not None:
    quinoa_builder_factory = waves.scons_extensions.sbatch_quinoa_builder_factory
envQuinoa.Append(
    BUILDERS={"QuinoaSolver": quinoa_builder_factory(program=envQuinoa["charmrun"], subcommand=envQuinoa["inciter"])}
)

# Comments used in tutorial code snippets: marker-6

# Add simulation targets
SConscript(
    "SConscript", variant_dir=env["variant_dir_base"], exports={"env": env, "envQuinoa": envQuinoa}, duplicate=True
)

# Comments used in tutorial code snippets: marker-7

# Add default target list to help message
env.Default()  # Empty defaults list to avoid building all simulation targets by default
# Add aliases to help message so users know what build target options are available
# This must come *after* all expected Alias definitions and SConscript files.
env.ProjectHelp()

# Comments used in tutorial code snippets: marker-8
