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

import ctypes
import enum
import json
import os
import random
import runpy
import shutil
import socket
import subprocess
import sys

from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, ArgumentTypeError, RawTextHelpFormatter
from tempfile import TemporaryDirectory
from types import SimpleNamespace

# ------------------------------------------------------------------------------
# Argument parser
# ------------------------------------------------------------------------------
prog = sys.argv[0]
epilog = f"Use '-h' to show command arguments; e.g. '{prog} new -h'"

parser = ArgumentParser(description="Salis: Simple A-Life Simulator", epilog=epilog, formatter_class=RawTextHelpFormatter, prog=prog)
sub_parsers = parser.add_subparsers(dest="command", required=True)
formatter_class = lambda prog: ArgumentDefaultsHelpFormatter(max_help_position=48, prog=prog)

new = sub_parsers.add_parser("new", formatter_class=formatter_class, help="create new simulation")
load = sub_parsers.add_parser("load", formatter_class=formatter_class, help="load saved simulation")
server = sub_parsers.add_parser("server", formatter_class=formatter_class, help="run data server")
client = sub_parsers.add_parser("client", formatter_class=formatter_class, help="run data client")

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

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

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

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

def port(i):
    ival = int(i, 0)
    if not 0 <= ival <= 65535: raise ArgumentTypeError("value must be valid port number")
    return ival

fmt_id = lambda val: val
fmt_hex = lambda val: hex(int(val, 0)) if type(val) == str else hex(val)

options = {
    (("A", "anc"), (new,), fmt_id): {"metavar": "ANC", "help": "ancestor file name without extension; to be compiled on all cores; ANC points to 'anc/{arch}/{ANC}.asm'", "required": True, "type": str},
    (("a", "arch"), (new,), fmt_id): {"choices": architectures, "help": "VM architecture", "default": "dummy", "required": False, "type": str},
    (("C", "clones"), (new,), fmt_id): {"metavar": "N", "help": "number of ancestor clones on each core", "default": 1, "required": False, "type": nat},
    (("c", "cores"), (new,), fmt_id): {"metavar": "N", "help": "number of simulator cores", "default": 2, "required": False, "type": nat},
    (("d", "data-push-pow"), (new,), fmt_id): {"metavar": "POW", "help": "data aggregation interval exponent; interval = 2^{POW} >= {sync-pow}; a value of 0 disables data aggregation; requires 'sqlite' and 'zlib'", "default": 28, "required": False, "type": pos},
    (("f", "force"), (new,), fmt_id): {"action": "store_true", "help": "overwrite existing simulation of given name", "required": False},
    (("F", "muta-flip"), (new,), fmt_id): {"action": "store_true", "help": "cosmic rays flip bits instead of randomizing whole bytes", "required": False},
    (("g", "c-compiler"), (new, load, server, client), fmt_id): {"metavar": "CC", "help": "C compiler to use", "default": "gcc", "required": False, "type": str},
    (("G", "c-compiler-flags"), (new, load, server, client), fmt_id): {"metavar": "FLAGS", "help": "base set of flags to pass to C compiler", "default": "-Wall -Wextra -Werror -pedantic", "required": False, "type": str},
    (("g++", "cpp-compiler"), (client,), fmt_id): {"metavar": "CXX", "help": "C++ compiler to use", "default": "g++", "required": False, "type": str},
    (("G++", "cpp-compiler-flags"), (client,), fmt_id): {"metavar": "FLAGS", "help": "base set of flags to pass to C++ compiler", "default": "-Wall -Wextra -Werror -pedantic", "required": False, "type": str},
    (("H", "home"), (new, load, server), fmt_id): {"metavar": "PATH", "help": "salis home directory", "default": os.path.join(os.environ["HOME"], ".salis"), "required": False, "type": str},
    (("i", "ip"), (client,), fmt_id): {"metavar": "IP", "help": "ip address of server", "default": "127.0.0.1", "required": False, "type": str},
    (("M", "muta-pow"), (new,), fmt_id): {"metavar": "POW", "help": "mutator range exponent; each step a cosmic ray hits addr, where addr = rand_uint64() %% 2^{POW}; lower values of POW mean higher mutation rates", "default": 32, "required": False, "type": pos},
    (("m", "mvec-pow"), (new,), fmt_id): {"metavar": "POW", "help": "memory core size exponent; size = 2^{POW}", "default": 20, "required": False, "type": pos},
    (("n", "name"), (new, load, server), fmt_id): {"metavar": "NAME", "help": "name of new or loaded simulation", "default": "def.sim", "required": False, "type": str},
    (("o", "optimized"), (new, load, server, client), fmt_id): {"action": "store_true", "help": "build with optimizations", "required": False},
    (("P", "port"), (server, client), fmt_id): {"metavar": "PORT", "help": "port number for data server", "default": 8080, "required": False, "type": port},
    (("p", "pre-cmd"), (new, load, server, client), fmt_id): {"metavar": "CMD", "help": "shell command with which to wrap call to executable; e.g. gdb, time, valgrind, etc.", "required": False, "type": str},
    (("s", "seed"), (new,), fmt_hex): {"metavar": "SEED", "help": "seed value for new simulation; a value of 0 disables cosmic rays; a value of -1 creates a random seed", "default": 0, "required": False, "type": seed},
    (("T", "keep-temp-dir"), (new, load, server, client), fmt_id): {"action": "store_true", "help": "keep temporary directory on exit", "required": False},
    (("t", "thread-gap"), (new, load), fmt_hex): {"metavar": "N", "help": "memory gap between core elements in bytes; may help reduce cache misses", "default": 0x100, "required": False, "type": nat},
    (("u", "ui"), (new, load), fmt_id): {"choices": uis, "help": "user interface", "default": "curses", "required": False, "type": str},
    (("x", "no-compress"), (new,), fmt_id): {"action": "store_true", "help": "do not compress save files; useful if 'zlib' is unavailable", "required": False},
    (("y", "sync-pow"), (new,), fmt_id): {"metavar": "POW", "help": "core sync interval exponent; sync events occur every N steps, where N = 2^{POW}", "default": 20, "required": False, "type": pos},
    (("z", "auto-save-pow"), (new,), fmt_id): {"metavar": "POW", "help": "auto-save interval exponent; auto-saves occur every N steps, where N = 2^{POW}", "default": 36, "required": False, "type": pos},
}

for (name, sub_parsers, _), kwargs in options.items():
    for sub_parser in sub_parsers:
        sub_parser.add_argument(f"-{name[0]}", f"--{name[1]}", **kwargs)

args = parser.parse_args()

# ------------------------------------------------------------------------------
# Build class
# ------------------------------------------------------------------------------
class Build:
    def __init__(self, path, log, library=False, cpp=False):
        self.log = log
        self.library = library

        self.tempdir = TemporaryDirectory(prefix="salis_", delete=not args.keep_temp_dir)
        self.log.info(f"Generated temporary directory for builds at: {self.tempdir.name}")

        self.name = os.path.splitext(os.path.basename(path))[0]
        self.binfile = os.path.join(self.tempdir.name, f"{self.name}{".so" if library else ""}")
        self.argsfile = os.path.join(self.tempdir.name, f"{self.name}.arg")

        self.flags = {*(args.cpp_compiler_flags if cpp else args.c_compiler_flags).split(), *({"-shared", "-fPIC"} if library else set()), *({"-O3"} if args.optimized else {"-ggdb"})}
        self.defines = {"-DNDEBUG"} if args.optimized else set()
        self.links = set()

        self.build_cmd = [args.cpp_compiler if cpp else args.c_compiler, f"@{self.argsfile}", path, "-o", self.binfile]
        self.log.info(f"Build class initialized for {"library" if library else "executable"}: {path}")
        self.log.info(f"Compiler flags stored at: {self.argsfile}")

    def build(self):
        fmt_nl = lambda line: f"{line}\n"
        fmt_define = lambda line: f"{line.replace(" ", "\\ ").replace("\"", "\\\"").replace("'", "\\'")}\n"

        with open(self.argsfile, "w") as f:
            f.writelines(map(fmt_nl, sorted(self.flags)))
            f.writelines(map(fmt_define, sorted(self.defines)))
            f.writelines(map(fmt_nl, sorted(self.links)))

        self.log.info(f"Running build command: '{self.build_cmd}'")
        subprocess.run(self.build_cmd, check=True)

    def exec(self):
        assert not self.library

        run_cmd = args.pre_cmd.split() if args.pre_cmd else []
        run_cmd.append(self.binfile)

        self.log.info(f"Running binary with command: '{" ".join(run_cmd)}'")
        proc = subprocess.Popen(run_cmd, stdout=sys.stdout, stderr=sys.stderr)

        # 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:
            proc.wait()
        except KeyboardInterrupt:
            proc.terminate()
            proc.wait()

        code = proc.returncode

        if code != 0: raise RuntimeError(f"Binary returned code: {code}")

# ------------------------------------------------------------------------------
# Logging
# ------------------------------------------------------------------------------
log_logs = []
log_ns = SimpleNamespace()
log_ns.info = lambda msg: log_logs.append(("info", msg))
log_ns.warn = lambda msg: log_logs.append(("warn", msg))
log_b = Build("core/logger.c", log_ns, library=True)
log_b.build()

# Pull in logging functions from C
# This way there's only a single unified logging system to care about
log = SimpleNamespace()
log_dll = ctypes.CDLL(log_b.binfile)
log.info = lambda msg: log_dll.log_info(msg.encode())
log.warn = lambda msg: log_dll.log_warn(msg.encode())

for log_level, log_msg in log_logs:
    getattr(log, log_level)(log_msg)

# ------------------------------------------------------------------------------
# Options dict formatter
# ------------------------------------------------------------------------------
fmt_h2u = lambda field: field.replace("-", "_")
fmt_u2h = lambda field: field.replace("_", "-")

def fmt_field(field, val):
    for ((_, name), _, fmt), _ in options.items():
        if name == fmt_u2h(field):
            return fmt(val)

    return fmt_id(val)

def fmt_opts(opts):
    opts_out = {}

    for field, val in opts.items():
        opts_out[field] = fmt_field(field, val)

    return json.dumps(opts_out, indent=2)

# ------------------------------------------------------------------------------
# Source configuration
# ------------------------------------------------------------------------------
log.info(f"Called '{prog} {args.command}' with the following options: {fmt_opts(vars(args))}")
ns = SimpleNamespace()

class Request(enum.Enum):
    REQUEST_NAME = "n"
    REQUEST_OPTS = "o"
    REQUEST_HASH = "h"

def gen_sim_paths():
    ns.sim_dir = os.path.join(args.home, args.name)
    ns.sim_opts = os.path.join(ns.sim_dir, "opts.json")
    ns.sim_path = os.path.join(ns.sim_dir, args.name)

def source_from_opts_file():
    if not os.path.isdir(ns.sim_dir):
        raise RuntimeError(f"No simulation found named {args.name}")

    with open(ns.sim_opts, "r") as f:
        opts = json.loads(f.read())

    for key, val in opts.items():
        setattr(args, key, val)

    log.info(f"Sourced configuration from '{ns.sim_opts}': {fmt_opts(opts)}")

def request_from_server(request):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client:
        client.connect((args.ip, args.port))
        client.sendall(request.value.encode())
        return json.load(client.makefile(mode="r"))

# New simulation
if args.command == "new":
    gen_sim_paths()

    if 0 < args.data_push_pow < args.sync_pow:
        raise RuntimeError("Data push power must be equal or greater than core sync power")

    if os.path.isdir(ns.sim_dir) and args.force:
        log.warn(f"Force flag used! Wiping old simulation at: {ns.sim_dir}")
        shutil.rmtree(ns.sim_dir)

    if os.path.isdir(ns.sim_dir):
        raise RuntimeError(f"Simulation directory found at: {ns.sim_dir}; use --force flag to remove it")

    if args.seed == -1:
        args.seed = random.getrandbits(64)
        log.info(f"Using random seed: {hex(args.seed)}")

    log.info(f"Creating new simulation directory at: {ns.sim_dir}")
    log.info(f"Creating configuration file at: {ns.sim_opts}")
    os.mkdir(ns.sim_dir)

    opts = {fmt_h2u(name[1]): getattr(args, fmt_h2u(name[1])) for (name, sub_parsers, _), kwargs in options.items() if new in sub_parsers and load not in sub_parsers}

    with open(ns.sim_opts, "w") as f:
        f.write(f"{fmt_opts(opts)}\n")

# Loaded simulation
if args.command == "load":
    gen_sim_paths()
    source_from_opts_file()

# Data server
if args.command == "server":
    gen_sim_paths()
    source_from_opts_file()

# Data client
if args.command == "client":
    # Pull basic information from the server
    # Required to build the client with correct flags
    args.name = request_from_server(Request.REQUEST_NAME)["name"]
    log.info(f"Sourced simulation name from '{args.ip}:{args.port}': {args.name}")

    opts = request_from_server(Request.REQUEST_OPTS)

    for key, val in opts.items():
        setattr(args, key, val)

    log.info(f"Sourced configuration from '{args.ip}:{args.port}': {fmt_opts(opts)}")

    # Server and client should be on the same hash
    server_hash = request_from_server(Request.REQUEST_HASH)["hash"]
    client_hash = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip()

    if server_hash != client_hash:
        raise RuntimeError(f"Server hash '{server_hash}' and client hash '{client_hash}' don't match")

    log.info(f"Confirmed server and client are running against the same git hash: {server_hash}")

# ------------------------------------------------------------------------------
# Architecture variables
# ------------------------------------------------------------------------------
arch_path = os.path.join("arch", args.arch, "vars.py")
log.info(f"Loading architecture variables from: {arch_path}")
arch_vars = runpy.run_path(arch_path, init_globals=globals())

# ------------------------------------------------------------------------------
# Assembler
# ------------------------------------------------------------------------------
anc_path = os.path.join("anc", args.arch, f"{args.anc}.asm")

if not os.path.isfile(anc_path):
    raise RuntimeError(f"Could not find ancestor file: {anc_path}")

with open(anc_path, "r") as f:
    lines = f.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)

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
            break

    if not found:
        raise RuntimeError(f"Unrecognized instruction in ancestor file: {line}")

anc_repr = f"{{{",".join(map(str, anc_bytes))}}}"
log.info(f"Compiled ancestor file '{anc_path}' into byte array: {anc_repr}")

# ------------------------------------------------------------------------------
# Compiler flags
# ------------------------------------------------------------------------------
def pop_ui_vars():
    ns.ui_path = os.path.join("ui", args.ui, "vars.py")
    ns.ui_vars = runpy.run_path(ns.ui_path, init_globals=globals())
    log.info(f"Sourced UI variables from: {ns.ui_path}")

    ns.b.flags.add(f"-Iui/{args.ui}")
    ns.b.flags.update({*ns.ui_vars["flags"]})
    ns.b.defines.update({*ns.ui_vars["defines"]})
    ns.b.links.update({*ns.ui_vars["links"]})

def pop_data_push_vars():
    if args.data_push_pow:
        ns.sim_db = os.path.join(ns.sim_dir, f"{args.name}.sqlite3")
        ns.b.defines.add(f"-DDATA_PUSH_PATH=\"{ns.sim_db}\"")
        ns.b.links.add("-lsqlite3")
        ns.b.links.add("-lz")
        log.info(f"Data will be aggregated at: {ns.sim_db}")
    else:
        log.warn("Data aggregation disabled")

    if not args.no_compress:
        ns.b.links.add("-lz")
        log.info("Save file compression enabled")
    else:
        log.warn("Save file compression disabled")

def pop_sim_path_vars():
    ns.b.defines.add(f"-DSIM_OPTS=\"{ns.sim_opts}\"")
    ns.b.defines.add(f"-DSIM_PATH=\"{ns.sim_path}\"")

def pop_net_vars():
    ns.b.defines.add(f"-DPORT={args.port}")
    ns.b.defines.add(f"-DPORT_STR=\"{args.port}\"")
    ns.b.defines.add(f"-DREQUEST_NAME='{Request.REQUEST_NAME.value}'")
    ns.b.defines.add(f"-DREQUEST_OPTS='{Request.REQUEST_OPTS.value}'")
    ns.b.defines.add(f"-DREQUEST_HASH='{Request.REQUEST_HASH.value}'")
    ns.b.links.add("-ljson-c")

def pop_general():
    ns.b.flags.add(f"-Iarch/{args.arch}")
    ns.b.flags.add("-Icore")

    if args.data_push_pow:
        ns.b.defines.add("-DDATA_PUSH")
        ns.b.defines.add(f"-DDATA_PUSH_INTERVAL={2 ** args.data_push_pow}ul")

    if not args.no_compress:
        ns.b.defines.add("-DCOMPRESS")

    if args.muta_flip:
        ns.b.defines.add("-DMUTA_FLIP")

    if arch_vars["mvec_loop"]:
        ns.b.defines.add("-DMVEC_LOOP")

    ns.b.defines.add(f"-DANC=\"{args.anc}\"")
    ns.b.defines.add(f"-DANC_BYTES={anc_repr}")
    ns.b.defines.add(f"-DANC_SIZE={len(anc_bytes)}")
    ns.b.defines.add(f"-DARCH=\"{args.arch}\"")
    ns.b.defines.add(f"-DAUTOSAVE_INTERVAL={2 ** args.auto_save_pow}ul")
    ns.b.defines.add(f"-DCLONES={args.clones}")
    ns.b.defines.add(f"-DCOMMAND_{args.command.upper()}")
    ns.b.defines.add(f"-DCORE_DATA_FIELDS={" ".join(f"CORE_DATA_FIELD({", ".join(field)})" for field in arch_vars["core_data_fields"])}")
    ns.b.defines.add(f"-DCORE_FIELD_COUNT={len(arch_vars["core_fields"])}")
    ns.b.defines.add(f"-DCORE_FIELDS={" ".join(f"CORE_FIELD({", ".join(field)})" for field in arch_vars["core_fields"])}")
    ns.b.defines.add(f"-DCORES={args.cores}")
    ns.b.defines.add(f"-DFOR_CORES={" ".join(f"FOR_CORE({i})" for i in range(args.cores))}")
    ns.b.defines.add(f"-DINST_COUNT={len(arch_vars["inst_set"])}")
    ns.b.defines.add(f"-DINST_SET={" ".join(f"INST({index}, {"_".join(inst[0])}, \"{" ".join(inst[0])}\", L'{inst[1]}')" for index, inst in enumerate(arch_vars["inst_set"]))}")
    ns.b.defines.add(f"-DMUTA_RANGE={2 ** args.muta_pow}ul")
    ns.b.defines.add(f"-DMVEC_SIZE={2 ** args.mvec_pow}ul")
    ns.b.defines.add(f"-DNAME=\"{args.name}\"")
    ns.b.defines.add(f"-DPROC_FIELD_COUNT={len(arch_vars["proc_fields"])}")
    ns.b.defines.add(f"-DPROC_FIELDS={" ".join(f"PROC_FIELD({", ".join(field)})" for field in arch_vars["proc_fields"])}")
    ns.b.defines.add(f"-DSEED={args.seed}ul")
    ns.b.defines.add(f"-DSYNC_INTERVAL={2 ** args.sync_pow}ul")

# Populate for new
if args.command == "new":
    ns.b = Build("core/salis.c", log)
    pop_ui_vars()
    pop_data_push_vars()
    pop_sim_path_vars()
    pop_general()

    ns.b.defines.add(f"-DAUTOSAVE_NAME_LEN={len(ns.sim_path) + 20}")
    ns.b.defines.add(f"-DTHREAD_GAP={args.thread_gap}")

# Populate for load
if args.command == "load":
    ns.b = Build("core/salis.c", log)
    pop_ui_vars()
    pop_data_push_vars()
    pop_sim_path_vars()
    pop_general()

    ns.b.defines.add(f"-DAUTOSAVE_NAME_LEN={len(ns.sim_path) + 20}")
    ns.b.defines.add(f"-DTHREAD_GAP={args.thread_gap}")

# Populate for server
if args.command == "server":
    ns.b = Build("data/server.c", log)
    pop_sim_path_vars()
    pop_net_vars()
    pop_general()

# Populate for client
if args.command == "client":
    ns.b = Build("data/client.cpp", log, cpp=True)
    pop_net_vars()
    pop_general()
    ns.b.defines.add(f"-DIP=\"{args.ip}\"")
    ns.b.links.add("-lGL")
    ns.b.links.add("-lglfw")
    ns.b.links.add("-limgui")
    ns.b.links.add("-limplot")
    ns.b.links.add("-lm")

# ------------------------------------------------------------------------------
# Build and launch executable
# ------------------------------------------------------------------------------
ns.b.build()
ns.b.exec()