# Released under the MIT License. See LICENSE for details.
#
"""Manage Ballistica execution environment.

This module is used to set up and/or check the global Python environment
before running a Ballistica app. This includes things such as paths,
logging, and app-dirs. Because these things are global in nature, this
should be done before any Ballistica modules are imported.

This module can also be exec'ed directly to set up a default environment
and then run the app.

Ballistica can be used without explicitly configuring the environment in
order to integrate it in arbitrary Python environments, but this may
cause some features to be disabled or behave differently than expected.
"""

import os
import sys
import time
import random
import logging
from pathlib import Path
from dataclasses import dataclass
from typing import TYPE_CHECKING
import __main__

if TYPE_CHECKING:
    from typing import Any, Callable

    from efro.logging import LogHandler

logger = logging.getLogger('ba.env')

# IMPORTANT - It is likely (and in some cases expected) that this
# module's code will be exec'ed multiple times. This is because it is
# the job of this module to set up Python paths for an engine run, and
# that may involve modifying sys.path in such a way that this module
# resolves to a different path afterwards (for example from
# /abs/path/to/ba_data/scripts/babase.py to ba_data/scripts/babase.py).
# This can result in the next import of baenv loading us from our 'new'
# location, which may or may not actually be the same file on disk as
# the last load. Either way, however, multiple execs can happen in some
# form.
#
# To handle that situation gracefully, we need to do a few things:
#
# - First, we need to store any mutable global state in the __main__
#   module; not in ourself. This way, alternate versions of ourself will
#   still know if we already ran configure/etc.
#
# - Second, we should avoid the use of isinstance and similar calls for
#   our types. An EnvConfig we create would technically be a different
#   type than an EnvConfig created by an alternate baenv.

# Build number and version of the ballistica binary we expect to be
# using.
TARGET_BALLISTICA_BUILD = 22915
TARGET_BALLISTICA_VERSION = '1.8.0a28'


@dataclass
class EnvConfig:
    """Final config values we provide to the engine."""

    #: Where app config/state data lives.
    config_dir: str

    #: Directory containing ba_data and any other platform-specific data.
    data_dir: str

    #: Where cache files live (files generated by the app which can be
    #: recreated if need be.
    cache_dir: str

    #: Where the app's built-in Python stuff lives.
    app_python_dir: str | None

    #: Where the app's built-in Python stuff lives in the default case.
    standard_app_python_dir: str

    #: Where the app's bundled third party Python stuff lives.
    site_python_dir: str | None

    #: Where custom Python provided by the user (mods) lives.
    user_python_dir: str | None

    #: We have a mechanism allowing :attr:`app_python_dir` to be
    #: overridden by placing a specially named directory in
    #: :attr:`user_python_dir`. This is true if that is enabled.
    is_user_app_python_dir: bool

    #: Our fancy app log handler. This handles feeding logs, stdout, and
    #: stderr into the engine so they show up on in-app consoles, etc.
    log_handler: LogHandler | None

    #: Initial data from the ``config.json`` file in the config dir.
    initial_app_config: Any

    #: Timestamp when we first started doing stuff.
    launch_time: float


@dataclass
class _EnvGlobals:
    """Globals related to baenv's operation.

    We store this in __main__ instead of in our own module because it is
    likely that multiple versions of our module will be spun up and we
    want a single set of globals (see notes at top of our module code).
    """

    config: EnvConfig | None = None
    called_configure: bool = False
    paths_set_failed: bool = False
    modular_main_called: bool = False

    @classmethod
    def get(cls) -> _EnvGlobals:
        """Create/return our singleton."""
        name = '_baenv_globals'
        envglobals: _EnvGlobals | None = getattr(__main__, name, None)
        if envglobals is None:
            envglobals = _EnvGlobals()
            setattr(__main__, name, envglobals)
        return envglobals


def did_paths_set_fail() -> bool:
    """Did we try to set paths and fail?

    :meta private:
    """
    return _EnvGlobals.get().paths_set_failed


def env_config_exists() -> bool:
    """Has a config been created?"""

    return _EnvGlobals.get().config is not None


def get_env_config() -> EnvConfig:
    """Return the active config, creating a default if none exists."""
    envglobals = _EnvGlobals.get()

    # If configure() has not been explicitly called, set up a
    # minimally-intrusive default config. We want Ballistica to default
    # to being a good citizen when imported into alien environments and
    # not blow away logging or otherwise muck with stuff. All official
    # paths to run Ballistica apps should be explicitly calling
    # configure() first to get a full featured setup.
    if not envglobals.called_configure:
        configure(setup_logging=False, setup_pycache_prefix=False)

    config = envglobals.config
    if config is None:
        raise RuntimeError(
            'baenv.configure() has been called but no config exists;'
            ' perhaps it errored?'
        )
    return config


def configure(
    *,
    config_dir: str | None = None,
    data_dir: str | None = None,
    user_python_dir: str | None = None,
    app_python_dir: str | None = None,
    site_python_dir: str | None = None,
    cache_dir: str | None = None,
    contains_python_dist: bool = False,
    setup_logging: bool = True,
    setup_pycache_prefix: bool = False,
    strict_threads_atexit: Callable[[Callable[[], None]], None] | None = None,
    launch_time: float | None = None,
) -> None:
    """Set up the environment for running a Ballistica app.

    This includes things such as Python path wrangling and app directory
    creation. This must be called before any actual Ballistica modules
    are imported; the environment is locked in as soon as that happens.

    ``launch_time`` is an optional epoch-seconds value (from
    :func:`time.time`) captured earlier in startup — typically at C++
    ``main()`` entry — used as the anchor for relative log timestamps.
    If not supplied, we fall back to sampling :func:`time.time` here.
    """

    # Prefer a caller-supplied launch_time (captured earlier in startup)
    # over our own sample so that log-timestamp 'relative time' reflects
    # real process start.
    if launch_time is None:
        launch_time = time.time()

    envglobals = _EnvGlobals.get()

    # Keep track of whether we've been *called*, not whether a config
    # has been created. Otherwise its possible to get multiple
    # overlapping configure calls going.
    if envglobals.called_configure:
        raise RuntimeError(
            'baenv.configure() has already been called;'
            ' it can only be called once.'
        )
    envglobals.called_configure = True

    # The very first thing we do is setup Python paths (while also
    # calculating some engine paths). This code needs to be bulletproof
    # since we have no logging yet at this point. We used to set up
    # logging first, but this way logging stuff will get loaded from its
    # proper final path (otherwise we might wind up using two different
    # versions of efro.logging in a single engine run).
    (
        user_python_dir,
        app_python_dir,
        site_python_dir,
        data_dir,
        config_dir,
        cache_dir,
        standard_app_python_dir,
        is_user_app_python_dir,
    ) = _setup_paths(
        user_python_dir,
        app_python_dir,
        site_python_dir,
        data_dir,
        config_dir,
        cache_dir,
    )

    # The one other thing we do before setting up logging is redirect
    # our pyc files to our cache dir. We want to do this is calced so
    # that as much stuff as possible (efro.logging), etc.) will get its
    # pyc files made in our custom cache dir.
    prev_pycache_prefix = sys.pycache_prefix
    if setup_pycache_prefix:
        sys.pycache_prefix = os.path.join(cache_dir, 'pyc')

    # Set up our log-handler and pipe Python's stdout/stderr into it.
    # Later, once the engine comes up, the handler will feed its logs
    # (including cached history) to the os-specific output location.
    # This means anything printed or logged at this point forward should
    # be visible on all platforms.
    log_handler = (
        _create_log_handler(launch_time, strict_threads_atexit)
        if setup_logging
        else None
    )

    # Load the raw app-config dict.
    app_config = _read_app_config(os.path.join(config_dir, 'config.json'))

    # Set logging levels to stored values or defaults.
    if setup_logging:
        _set_log_levels(app_config)

    # We want to always be run in UTF-8 mode; complain if we're not.
    if sys.flags.utf8_mode != 1:
        logger.warning(
            "Python's UTF-8 mode is not set. Running Ballistica without"
            ' it may lead to errors.'
        )

    # We (possibly) set pycache_prefix above so that opt .pyc files are
    # written to the cache directory that we just set up, but ideally
    # Python should have been set to that value at startup so that
    # modules we've imported up to this point get cached there too.
    #
    # In most cases we can actually do this by calcing/setting the same
    # path we use here before spinning up Python, but in some cases
    # that's impossible (such as our _modular_main path below where we
    # are already in Python before we get a chance to parse args that
    # affect cache path).
    #
    # So let's warn here any time we're trying to set up pycache_prefix
    # but find that we're setting it to a different value than it was
    # already set to. We can inform the user (or ourselves) how to line
    # things up using PYTHONPYCACHEPREFIX or whatnot.
    if setup_pycache_prefix and prev_pycache_prefix != sys.pycache_prefix:
        logger.warning(
            'Changing sys.pycache_prefix from %s to %s.'
            ' For best performance, run with PYTHONPYCACHEPREFIX=%s.',
            repr(prev_pycache_prefix),
            repr(sys.pycache_prefix),
            repr(sys.pycache_prefix),
        )

    # Attempt to create dirs that we'll write stuff to.
    _setup_dirs(config_dir, user_python_dir, cache_dir)

    # In debug builds, if we've not imported engine stuff yet, Kill off
    # random cache files occasionally to help ensure that code responds
    # correctly if/when the OS does the same thing.
    if __debug__:
        if '_babase' not in sys.modules:
            _cache_ninja_rampage(cache_dir)

    # Get ssl working if needed so we can use https and all that.
    _setup_certs(contains_python_dist)

    # This is now the active config.
    envglobals.config = EnvConfig(
        config_dir=config_dir,
        data_dir=data_dir,
        cache_dir=cache_dir,
        user_python_dir=user_python_dir,
        app_python_dir=app_python_dir,
        standard_app_python_dir=standard_app_python_dir,
        site_python_dir=site_python_dir,
        log_handler=log_handler,
        is_user_app_python_dir=is_user_app_python_dir,
        initial_app_config=app_config,
        launch_time=launch_time,
    )


def _cache_ninja_rampage(cache_dir: str) -> None:
    assert os.path.isdir(cache_dir)
    for basename, _dirnames, filenames in os.walk(cache_dir):
        for fname in filenames:
            # Let's kill one out of every 1000 files; should be a
            # reasonable amount of chaos I think. Can recalibrate this
            # as our average cache file count goes up.
            if random.random() < 0.001:
                fullpath = os.path.join(basename, fname)
                logging.getLogger('ba.cache').debug(
                    "Cache-ninja assassinated '%s'.", fullpath
                )
                # The whole point of this feature is that downstream
                # code must handle missing cache files; the kill itself
                # is not load-bearing. Swallow OSError so a read-only
                # cache (sandboxed envs, restrictive perms) or a file
                # that vanished mid-walk doesn't fatal startup.
                try:
                    os.unlink(fullpath)
                except OSError:
                    pass


def _read_app_config(config_file_path: str) -> dict:
    """Read the app config."""
    import json

    config: dict | Any
    config_contents = ''
    try:
        if os.path.exists(config_file_path):
            with open(config_file_path, encoding='utf-8') as infile:
                config_contents = infile.read()
            config = json.loads(config_contents)
            if not isinstance(config, dict):
                raise RuntimeError('Got non-dict for config root.')
        else:
            config = {}

    except Exception:
        logger.exception(
            "Error reading config file '%s'.\n"
            "Backing up broken config to'%s.broken'.",
            config_file_path,
            config_file_path,
        )

        try:
            import shutil

            shutil.copyfile(config_file_path, config_file_path + '.broken')
        except Exception:
            logger.exception('Error copying broken config.')
        config = {}

    return config


def _calc_data_dir(data_dir: str | None) -> str:
    if data_dir is None:
        # To calc default data_dir, we assume this module was imported
        # from that dir's ba_data/python subdir.
        assert Path(__file__).parts[-3:-1] == ('ba_data', 'python')
        data_dir_path = Path(__file__).parents[2]

        # Prefer tidy relative paths like './ba_data' if possible so
        # that things like stack traces are easier to read. For best
        # results, platforms where CWD doesn't matter can chdir to where
        # ba_data lives before calling configure().
        #
        # NOTE: If there's ever a case where the user is chdir'ing at
        # runtime we might want an option to use only abs paths here.
        cwd_path = Path.cwd()
        data_dir = str(
            data_dir_path.relative_to(cwd_path)
            if data_dir_path.is_relative_to(cwd_path)
            else data_dir_path
        )
    return data_dir


def _create_log_handler(
    launch_time: float,
    strict_threads_atexit: Callable[[Callable[[], None]], None] | None,
) -> LogHandler:
    from efro.logging import setup_logging, LogLevel

    log_handler = setup_logging(
        log_path=None,
        level=LogLevel.INFO,
        log_stdout_stderr=True,
        cache_size_limit=1024 * 1024,
        launch_time=launch_time,
        strict_threads=strict_threads_atexit is not None,
    )

    # If we were given a strict_threads_atexit call, it means we should
    # NOT use daemon threads but instead can use the atexit call to
    # register a callback to gracefully exit our handler thread just
    # before the interpreter shuts down. This is safer than using daemon
    # threads, which can theoretically continue to use Python objs
    # during and after interpreter shutdown.
    if strict_threads_atexit is not None:
        strict_threads_atexit(log_handler.shutdown)

    return log_handler


def _set_log_levels(app_config: dict) -> None:

    from bacommon.logging import get_base_logger_control_config_client
    from bacommon.loggercontrol import LoggerControlConfig

    try:
        # If BA_LOG_LEVELS env var is set, it completely overrides any
        # stored config. Format: 'logger=LEVEL,logger=LEVEL,...'
        # Example: 'ba.net=DEBUG,ba.connectivity=DEBUG'
        env_log_levels = os.environ.get('BA_LOG_LEVELS')
        if env_log_levels is not None:
            _apply_env_log_levels(env_log_levels)
            return

        config = app_config.get('Log Levels', None)

        if config is None:
            get_base_logger_control_config_client().apply()
            return

        # Make sure data is expected types/values since this is user
        # editable.
        valid_levels = {
            logging.NOTSET,
            logging.DEBUG,
            logging.INFO,
            logging.WARNING,
            logging.ERROR,
            logging.CRITICAL,
        }
        for logname, loglevel in config.items():
            if (
                not isinstance(logname, str)
                or not logname
                or not isinstance(loglevel, int)
                or not loglevel in valid_levels
            ):
                raise ValueError("Invalid 'Log Levels' data read from config.")

        get_base_logger_control_config_client().apply_diff(
            LoggerControlConfig(levels=config)
        ).apply()

    except Exception:
        logger.exception('Error setting log levels.')


def _apply_env_log_levels(env_val: str) -> None:
    """Apply log levels from the BA_LOG_LEVELS env var.

    Completely overrides stored config. Base defaults are applied
    first, then env var values are layered on top.
    """
    from bacommon.logging import get_base_logger_control_config_client
    from bacommon.loggercontrol import LoggerControlConfig

    level_names = {
        'NOTSET': logging.NOTSET,
        'DEBUG': logging.DEBUG,
        'INFO': logging.INFO,
        'WARNING': logging.WARNING,
        'ERROR': logging.ERROR,
        'CRITICAL': logging.CRITICAL,
    }

    levels: dict[str, int] = {}
    for entry in env_val.split(','):
        entry = entry.strip()
        if not entry:
            continue
        parts = entry.split('=', 1)
        if len(parts) != 2:
            raise ValueError(
                f'Invalid BA_LOG_LEVELS entry: {entry!r}'
                ' (expected logger=LEVEL)'
            )
        logname, levelstr = parts[0].strip(), parts[1].strip().upper()
        if levelstr not in level_names:
            valid = ', '.join(level_names)
            raise ValueError(
                f'Invalid log level {levelstr!r} in BA_LOG_LEVELS'
                f' (expected one of {valid})'
            )
        levels[logname] = level_names[levelstr]

    get_base_logger_control_config_client().apply_diff(
        LoggerControlConfig(levels=levels)
    ).apply()


def _setup_certs(contains_python_dist: bool) -> None:
    # In situations where we're bringing our own Python, let's also
    # provide our own root certs so ssl works. We can consider
    # overriding this in particular embedded cases if we can verify that
    # system certs are working. We also allow forcing this via an env
    # var if the user desires.
    if (
        contains_python_dist
        or os.environ.get('BA_USE_BUNDLED_ROOT_CERTS') == '1'
    ):
        import certifi

        # Let both OpenSSL and requests (if present) know to use this.
        os.environ['SSL_CERT_FILE'] = os.environ['REQUESTS_CA_BUNDLE'] = (
            certifi.where()
        )


def _setup_paths(
    user_python_dir: str | None,
    app_python_dir: str | None,
    site_python_dir: str | None,
    data_dir: str | None,
    config_dir: str | None,
    cache_dir: str | None,
) -> tuple[str | None, str | None, str | None, str, str, str, str, bool]:
    # pylint: disable=too-many-positional-arguments

    # First a few paths we can ALWAYS calculate since they don't affect
    # Python imports:

    envglobals = _EnvGlobals.get()

    data_dir = _calc_data_dir(data_dir)

    # Default config-dir is simply ~/.bombsquad
    if config_dir is None:
        config_dir = str(Path(Path.home(), '.bombsquad'))

    # By default, cache-dir is simply 'cache' under config-dir.
    if cache_dir is None:
        cache_dir = str(Path(config_dir, 'cache'))

    # Standard app-python-dir is simply ba_data/python under data-dir.
    standard_app_python_dir = str(Path(data_dir, 'ba_data', 'python'))

    # Whether the final app-dir we're returning is a custom user-owned one.
    is_user_app_python_dir = False

    # If _babase has already been imported, there's not much we can do
    # at this point aside from complain and inform for next time.
    if '_babase' in sys.modules:
        app_python_dir = user_python_dir = site_python_dir = None

        # We don't actually complain yet here; we simply take note that
        # we weren't able to set paths. Then we complain if/when the app
        # is started. This way, non-app uses of babase won't be filled
        # with unnecessary warnings.
        envglobals.paths_set_failed = True

    else:
        # Ok; _babase hasn't been imported yet, so we can muck with
        # Python paths.

        if app_python_dir is None:
            app_python_dir = standard_app_python_dir

        # Likewise site-python-dir defaults to
        # ba_data/python-site-packages.
        if site_python_dir is None:
            site_python_dir = str(
                Path(data_dir, 'ba_data', 'python-site-packages')
            )

        # By default, user-python-dir is simply 'mods' under config-dir.
        if user_python_dir is None:
            user_python_dir = str(Path(config_dir, 'mods'))

        # Wherever our user_python_dir is, if we find a sys/FOO_BAR dir
        # under it where FOO matches our version and BAR matches our
        # build number, use that as our app_python_dir. This allows
        # modding built-in stuff on platforms where there is no write
        # access to said built-in stuff.
        check_dir = Path(
            user_python_dir,
            'sys',
            f'{TARGET_BALLISTICA_VERSION}_{TARGET_BALLISTICA_BUILD}',
        )
        try:
            if check_dir.is_dir():
                app_python_dir = str(check_dir)
                is_user_app_python_dir = True
        except PermissionError:
            logger.warning(
                "PermissionError checking user-app-python-dir path '%s'.",
                check_dir,
            )

        # Ok, now apply these to sys.path.

        # First off, strip out any instances of the path containing this
        # module. We will *probably* be re-adding the same path in a
        # moment so this keeps things cleaner. Though hmm should we
        # leave it in there in cases where we *don't* re-add the same
        # path?...
        our_parent_path = Path(__file__).parent.resolve()
        oldpaths: list[str] = [
            p for p in sys.path if Path(p).resolve() != our_parent_path
        ]

        # Let's place mods first (so users can override whatever they
        # want) followed by our app scripts and lastly our bundled site
        # stuff.

        # One could make the argument that at least our bundled app &
        # site stuff should be placed at the end so actual local site
        # stuff could override it. That could be a good thing or a bad
        # thing. Maybe we could add an option for that, but for now I'm
        # prioritizing our stuff to give as consistent an environment as
        # possible.
        ourpaths = [user_python_dir, app_python_dir, site_python_dir]

        # Special case: our modular builds will have a 'python-dylib'
        # dir alongside the 'python' scripts dir which contains our
        # binary Python modules. If we see that, add it to the path also.
        # Not sure if we'd ever have a need to customize this path.
        dylibdir = f'{app_python_dir}-dylib'
        if os.path.exists(dylibdir):
            ourpaths.append(dylibdir)

        sys.path = ourpaths + oldpaths

    return (
        user_python_dir,
        app_python_dir,
        site_python_dir,
        data_dir,
        config_dir,
        cache_dir,
        standard_app_python_dir,
        is_user_app_python_dir,
    )


def _setup_dirs(
    config_dir: str | None, user_python_dir: str | None, cache_dir: str
) -> None:
    create_dirs: list[tuple[str, str | None]] = [
        ('config', config_dir),
        ('cache', cache_dir),
        ('user_python', user_python_dir),
    ]
    for cdirname, cdir in create_dirs:
        if cdir is not None:
            try:
                os.makedirs(cdir, exist_ok=True)
            except Exception:
                # Not the end of the world if we can't make these dirs.
                logger.warning(
                    "Unable to create %s dir at '%s'.", cdirname, cdir
                )


def _extract_arg(args: list[str], names: list[str], is_dir: bool) -> str | None:
    """Given a list of args and an arg name, returns a value.

    The arg flag and value are removed from the arg list. We also check
    to make sure the path exists.

    raises CleanErrors on any problems.
    """
    from efro.error import CleanError

    count = sum(args.count(n) for n in names)
    if not count:
        return None

    if count > 1:
        raise CleanError(f'Arg {names} passed multiple times.')

    for name in names:
        if name not in args:
            continue
        argindex = args.index(name)
        if argindex + 1 >= len(args):
            raise CleanError(f'No value passed after {name} arg.')

        val = args[argindex + 1]
        del args[argindex : argindex + 2]

        if is_dir and not os.path.isdir(val):
            namepretty = names[0].removeprefix('--')
            raise CleanError(
                f"Provided {namepretty} path '{val}' is not a directory."
            )
        return val

    raise RuntimeError(f'Expected arg name not found from {names}')


def _modular_main() -> None:
    from efro.error import CleanError

    # Fundamentally, running a Ballistica app consists of the following:
    # import baenv; baenv.configure(); import babase; babase.app.run()
    #
    # First baenv sets up things like Python paths the way the engine
    # needs them, and then we import and run the engine.
    #
    # Below we're doing a slightly fancier version of that. Namely, we
    # do some processing of command line args to allow overriding of
    # paths or running explicit commands or whatever else. Our goal is
    # that this modular form of the app should be basically
    # indistinguishable from the monolithic form when used from the
    # command line.

    try:

        # Take note that we're running via modular-main. The native
        # layer can key off this to know whether it should apply
        # sys.argv or not.
        _EnvGlobals.get().modular_main_called = True

        # Deal with a few key things here ourself before even running
        # configure.

        # The extract_arg stuff below modifies this so we work with a
        # copy.
        args = sys.argv.copy()

        # NOTE: We need to keep these arg long/short arg versions synced
        # to those in core_config.cc. That code will parse these same
        # args (even if it doesn't do anything with them in this modular
        # path) and will complain if unrecognized args come through.

        # Our -c arg basically mirrors Python's -c arg. If we get that,
        # simply exec it and return; no engine stuff.
        command = _extract_arg(args, ['--command', '-c'], is_dir=False)
        if command is not None:
            exec(command)  # pylint: disable=exec-used
            return

        config_dir = _extract_arg(args, ['--config-dir', '-C'], is_dir=True)
        data_dir = _extract_arg(args, ['--data-dir', '-d'], is_dir=True)
        mods_dir = _extract_arg(args, ['--mods-dir', '-m'], is_dir=True)
        cache_dir = _extract_arg(args, ['--cache-dir', '-a'], is_dir=True)

        # We run configure() BEFORE importing babase. (part of its job
        # is to wrangle paths which can affect where babase and
        # everything else gets loaded from).
        configure(
            config_dir=config_dir,
            data_dir=data_dir,
            user_python_dir=mods_dir,
            cache_dir=cache_dir,
        )

        import babase

        # The engine will have parsed and processed all other args as
        # part of the above import. If there were errors or args such as
        # --help which should lead to us immediately returning, do so.
        code = babase.get_immediate_return_code()
        if code is not None:
            sys.exit(code)

        # Aaaand we're off!
        babase.app.run()

    # Code wanting us to die with a clean error message instead of an
    # ugly stack trace can raise one of these.
    except CleanError as clean_exc:
        clean_exc.pretty_print()
        sys.exit(1)


# Exec'ing this module directly will do a standard app run.
if __name__ == '__main__':
    _modular_main()
