#!/usr/bin/env -S PYTHONDONTWRITEBYTECODE=1 python3

# Author:  Paul Oliver <contact@pauloliver.dev>
# Project: salis-v3

# Salis simulator launcher script
# Emits a single C source file, builds it into a binary and launches it.
# JIT compilation allows quick switching between all available executable configurations.

import os
import random
import shutil
import subprocess
import sys

from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, ArgumentTypeError, RawTextHelpFormatter
from jinja2 import Environment, FileSystemLoader, StrictUndefined
from tempfile import TemporaryDirectory

# ------------------------------------------------------------------------------
# Parse CLI arguments
# ------------------------------------------------------------------------------
headline = "Salis: Simple A-Life Simulator"
script   = sys.argv[0]
epilog   = f"Use '-h' to list arguments for each command.\nExample: '{script} bench -h'"

main_parser = ArgumentParser(
    description     = headline,
    epilog          = epilog,
    formatter_class = RawTextHelpFormatter,
    prog            = script,
)

parsers = main_parser.add_subparsers(dest="command", required=True)
fclass  = ArgumentDefaultsHelpFormatter

bench = parsers.add_parser("bench", formatter_class=fclass, help="run benchmark")
load  = parsers.add_parser("load", formatter_class=fclass, help="load saved simulation")
new   = parsers.add_parser("new", formatter_class=fclass, help="create new simulation")

archs = os.listdir("./arch")
uis   = os.listdir("./ui")

def iseed(i):
    ival = int(i, 0)
    if ival < -1: raise ArgumentTypeError("invalid seed value")
    return ival

def ipos(i):
    ival = int(i, 0)
    if ival < 0: raise ArgumentTypeError("value must be positive integer")
    return ival

def inat(i):
    ival = int(i, 0)
    if ival < 1: raise ArgumentTypeError("value must be greater than zero")
    return ival

option_keys = ["short", "long", "metavar", "description", "default", "required", "type", "parsers"]

# fmt: off
option_conf = [
    ["A", "anc",             "ANC",  "ancestor file name without extension, to be compiled on "
                                     "all cores (ANC points to 'ancs/<ARCH>/<ANC>.asm')",        None,      True,  str,   [bench, new]],
    ["a", "arch",            archs,  "VM architecture",                                          "dummy",   False, str,   [bench, new]],
    ["b", "steps",           "N",    "number of steps to run in benchmark",                      0x1000000, False, ipos,  [bench]],
    ["C", "clones",          "N",    "number of ancestor clones on each core",                   1,         False, inat,  [bench, new]],
    ["c", "cores",           "N",    "number of simulator cores",                                2,         False, inat,  [bench, new]],
    ["d", "data-push-pow",   "POW",  "data aggregation interval exponent (interval == 2^POW >= "
                                     "thread sync interval); a value of 0 disables data "
                                     "aggregation (requires 'sqlite')",                          28,        False, ipos,  [new]],
    ["f", "force",           None,   "overwrite existing simulation of given name",              False,     False, bool,  [new]],
    ["F", "muta-flip",       None,   "cosmic rays flip bits instead of randomizing whole bytes", False,     False, bool,  [bench, new]],
    ["M", "muta-pow",        "POW",  "mutator range exponent (range == 2^POW)",                  32,        False, ipos,  [bench, new]],
    ["m", "mvec-pow",        "POW",  "memory vector size exponent (size == 2^POW)",              20,        False, ipos,  [bench, new]],
    ["n", "name",            "NAME", "name of new or loaded simulation",                         "def.sim", False, str,   [load, new]],
    ["o", "optimized",       None,   "builds salis binary with optimizations",                   False,     False, bool,  [bench, load, new]],
    ["p", "pre-cmd",         "CMD",  "shell command to wrap call to executable (e.g. gdb, "
                                     "valgrind, etc.)",                                          None,      False, str,   [bench, load, new]],
    ["s", "seed",            "SEED", "seed value for new simulation; a value of 0 disables "
                                     "cosmic rays; a value of -1 creates a random seed",         0,         False, iseed, [bench, new]],
    ["S", "print-source",    None,   "print generated C source to stdout and exit",              False,     False, bool,  [bench, load, new]],
    ["T", "delete-temp-dir", None,   "delete temporary directory on exit",                       True,      False, bool,  [bench, load, new]],
    ["t", "thread-gap",      "N",    "memory gap between cores in bytes (may help reduce cache "
                                     "misses?)",                                                 0x100,     False, inat,  [bench, load, new]],
    ["u", "ui",              uis,    "user interface",                                           "curses",  False, str,   [load, new]],
    ["x", "compress",        None,   "compress save files (requires 'zlib')",                    True,      False, bool,  [new]],
    ["y", "sync-pow",        "POW",  "core sync interval exponent (interval == 2^POW)",          20,        False, ipos,  [bench, new]],
    ["z", "auto-save-pow",   "POW",  "auto-save interval exponent (interval == 2^POW)",          36,        False, ipos,  [new]],
]
# fmt: on

# Map arguments to subparsers that use them
options    = list(map(lambda option: dict(zip(option_keys, option)), option_conf))
parser_map = ((parser, option) for option in options for parser in option["parsers"])

for parser, option in parser_map:
    arg_kwargs = {}

    def push_same(key):
        arg_kwargs[key] = option[key]

    def push_diff(tgt_key, src_key):
        arg_kwargs[tgt_key] = option[src_key]

    def push_val(key, val):
        arg_kwargs[key] = val

    push_diff("help", "description")
    push_same("required")

    # No metavar means this argument is a flag
    if option["metavar"] is None:
        push_val("action", "store_false" if option["default"] else "store_true")
    else:
        push_same("default")
        push_same("type")

    if type(option["metavar"]) is list:
        push_diff("choices", "metavar")

    if type(option["metavar"]) is str:
        push_same("metavar")

    parser.add_argument(
        f"-{option["short"]}",
        f"--{option["long"]}",
        **arg_kwargs,
    )

args = main_parser.parse_args()

def info(msg, val=""):
    print(f"\033[1;34mINFO:\033[0m {msg}", val)

def warn(msg, val=""):
    print(f"\033[1;31mWARN:\033[0m {msg}", val)

def error(msg, val=""):
    print(f"\033[1;31mERROR:\033[0m {msg}", val)
    sys.exit(1)

# ------------------------------------------------------------------------------
# Load configuration
# ------------------------------------------------------------------------------
info(headline)
info(f"Called '{script}' with the following options:")

for key, val in vars(args).items():
    print(f"{key} = {repr(val)}")

if args.command in ["load", "new"]:
    sim_dir  = f"{os.environ["HOME"]}/.salis/{args.name}"
    sim_opts = f"{sim_dir}/opts.py"
    sim_path = f"{sim_dir}/{args.name}"

if args.command in ["load"]:
    if not os.path.isdir(sim_dir):
        error("No simulation found named:", args.name)

    info(f"Sourcing configuration from '{sim_opts}':")
    sys.path.append(sim_dir)
    import opts as opts_module

    # Copy all fields in configuration file into the 'args' object
    opts = (opt for opt in dir(opts_module) if not opt.startswith("__"))

    for opt in opts:
        opt_attr = getattr(opts_module, opt)
        print(f"{opt} = {repr(opt_attr)}")
        setattr(args, opt, opt_attr)

if args.command in ["new"]:
    if args.data_push_pow != 0 and args.data_push_pow < args.sync_pow:
        error("Data push power must be equal or greater than thread sync power")

    if os.path.isdir(sim_dir) and args.force:
        warn("Force flag used - wiping old simulation at:", sim_dir)
        shutil.rmtree(sim_dir)

    if os.path.isdir(sim_dir):
        error("Simulation directory found at:", sim_dir)

    info("Creating new simulation directory at:", sim_dir)
    os.mkdir(sim_dir)

    info("Creating configuration file at:", sim_opts)

    opts = (
        option["long"].replace("-", "_")
        for option in options
        if new in option["parsers"] and load not in option["parsers"]
    )

    with open(sim_opts, "w") as file:
        for opt in opts:
            file.write(f"{opt} = {repr(eval(f"args.{opt}"))}\n")

# ------------------------------------------------------------------------------
# Load architecture and UI variables
# ------------------------------------------------------------------------------
arch_path = f"arch/{args.arch}"
info("Loading architecture specific variables from:", f"{arch_path}/arch_vars.py")
sys.path.append(arch_path)
import arch_vars

if args.command in ["load", "new"]:
    ui_path = f"ui/{args.ui}"
    info("Loading UI specific variables from:", f"{ui_path}/ui_vars.py")
    sys.path.append(ui_path)
    import ui_vars

# ------------------------------------------------------------------------------
# Fill in template variables
# ------------------------------------------------------------------------------
ul_val = lambda val: f"{hex(val)}ul"
ul_pow = lambda val: f"{hex(2 ** val)}ul"

includes = [
    "assert.h",
    "stdarg.h",
    "stdbool.h",
    "stddef.h",
    "stdint.h",
    "stdlib.h",
    "string.h",
    "threads.h",
]

inst_cap      = "0x80"
inst_mask     = "0x7f"
ipc_flag      = "0x80"
mall_flag     = "0x80"
muta_range    = ul_pow(args.muta_pow)
mvec_size     = ul_pow(args.mvec_pow)
sync_interval = ul_pow(args.sync_pow)
thread_gap    = ul_val(args.thread_gap)
uint64_half   = ul_val(0x8000000000000000)

if args.seed == -1:
    args.seed = ul_val(random.getrandbits(64))
    info("Using random seed", args.seed)
else:
    args.seed = ul_val(args.seed)

if args.command in ["bench"]:
    includes.append("stdio.h")
    args.steps = ul_val(args.steps)

if args.command in ["load", "new"]:
    auto_save_interval = ul_pow(args.auto_save_pow)
    auto_save_name_len = f"{len(sim_path) + 20}"

    if args.data_push_pow != 0:
        data_push_path         = f"{sim_dir}/{args.name}.sqlite3"
        data_push_interval     = ul_pow(args.data_push_pow)
        data_push_busy_timeout = 600000
        includes.append("sqlite3.h")
        info("Data will be aggregated at:", data_push_path)
    else:
        warn("Data aggregation disabled")

    if args.compress:
        includes.append("zlib.h")
        info("Save file compression enabled")
    else:
        warn("Save file compression disabled")

    includes.extend(ui_vars.includes)

# ------------------------------------------------------------------------------
# Assemble ancestor organism into byte array
# ------------------------------------------------------------------------------
if args.command in ["bench", "new"] and args.anc is not None:
    anc_path = f"ancs/{args.arch}/{args.anc}.asm"

    if not os.path.isfile(anc_path):
        error("Could not find ancestor file:", anc_path)

    with open(anc_path, "r") as file:
        lines = file.read().splitlines()

    lines = filter(lambda line: not line.startswith(";"), lines)
    lines = filter(lambda line: not line.isspace(), lines)
    lines = filter(lambda line: line, lines)
    lines = map(lambda line: line.split(), lines)

    # A very simple assembler that compares lines in input ASM file against
    # all entries in the instruction set table provided by each architecture.
    # The resulting bytes equate to each instruction's index on the table.
    anc_bytes = []

    for line in lines:
        found = False

        for byte, tup in enumerate(arch_vars.inst_set):
            if line == tup[0]:
                anc_bytes.append(byte)
                found = True
                continue

        if not found:
            error("Unrecognized instruction in ancestor file:", line)

    info(f"Compiled ancestor file '{anc_path}' into byte array:", anc_bytes)

# ------------------------------------------------------------------------------
# Emit C source
# ------------------------------------------------------------------------------
tempdir = TemporaryDirectory(prefix="salis_", delete=args.delete_temp_dir)
info("Created a temporary salis directory at:", tempdir.name)

salis_src = f"{tempdir.name}/salis.c"
info("Emitting C source at:", salis_src)

jinja_env = Environment(
    loader        = FileSystemLoader("."),
    lstrip_blocks = True,
    trim_blocks   = True,
    undefined     = StrictUndefined,
)

source_str = jinja_env.get_template("core.j2.c").render(**locals())

if args.print_source:
    info("Printing C source and exiting...")
    print(source_str)
    exit(0)

with open(salis_src, "w") as file:
    file.write(source_str)

# ------------------------------------------------------------------------------
# Build executable
# ------------------------------------------------------------------------------
salis_bin = f"{tempdir.name}/salis_bin"
info("Building salis binary at:", salis_bin)

build_cmd = ["gcc", salis_src, "-o", salis_bin, "-Wall", "-Wextra", "-Werror", "-Wno-overlength-strings", "-pedantic", "-std=c11"]
build_cmd.extend(["-O3", "-DNDEBUG"] if args.optimized else ["-ggdb"])

if args.command in ["load", "new"]:
    build_cmd.extend(ui_vars.flags)

    # Enable POSIX extensions (open_memstream)
    build_cmd.extend(["-lz", "-D_POSIX_C_SOURCE=200809L"] if args.compress else [])

    # Enable GNU extensions (asprintf)
    # This allows managing large SQL strings more easily
    build_cmd.extend(["-lsqlite3", "-D_GNU_SOURCE"] if args.data_push_pow != 0 else [])

info("Using build command:", " ".join(build_cmd))
subprocess.run(build_cmd, check=True)

# ------------------------------------------------------------------------------
# Run salis binary
# ------------------------------------------------------------------------------
info("Running salis binary...")

run_cmd = [args.pre_cmd] if args.pre_cmd else []
run_cmd.append(salis_bin)

info("Using run command:", " ".join(run_cmd))
salis_sp = subprocess.Popen(run_cmd, stdout=sys.stdout, stderr=sys.stderr)

# Ctrl-C terminates the simulator gracefully.
# When using signals (e.g. SIGTERM), they must be sent to the entire process group
# to make sure both the simulator and the interpreter get shut down.
try:
    salis_sp.wait()
except KeyboardInterrupt:
    salis_sp.terminate()
    salis_sp.wait()

code = salis_sp.returncode

if code != 0:
    error("Salis binary returned code:", code)