"""Tools for caching xonsh code."""

import hashlib
import marshal
import os
import sys

from xonsh import __version__ as XONSH_VERSION
from xonsh.built_ins import XSH
from xonsh.lib.lazyasd import lazyobject
from xonsh.platform import PYTHON_VERSION_INFO_BYTES
from xonsh.tools import is_writable_file, print_warning


def _splitpath(path, sofar=()):
    folder, path = os.path.split(path)
    if path == "":
        return sofar[::-1]
    elif folder == "":
        return (sofar + (path,))[::-1]
    else:
        return _splitpath(folder, sofar + (path,))


@lazyobject
def _CHARACTER_MAP():
    cmap = {chr(o): "_" + chr(o + 32) for o in range(65, 91)}
    cmap.update({".": "_.", "_": "__"})
    return cmap


def _cache_renamer(path, code=False):
    if not code:
        path = os.path.realpath(path)
    o = ["".join(_CHARACTER_MAP.get(i, i) for i in w) for w in _splitpath(path)]
    o[-1] = f"{o[-1]}.{sys.implementation.cache_tag}"
    return o


def should_use_cache(execer, mode):
    """
    Return ``True`` if caching has been enabled for this mode (through command
    line flags or environment variables)
    """
    if mode == "exec":
        return (execer.scriptcache or execer.cacheall) and (
            XSH.env["XONSH_CACHE_SCRIPTS"] or XSH.env["XONSH_CACHE_EVERYTHING"]
        )
    else:
        return execer.cacheall or XSH.env["XONSH_CACHE_EVERYTHING"]


def run_compiled_code(code, glb, loc, mode):
    """
    Helper to run code in a given mode and context.
    Returns a sys.exc_info() triplet in case the code raises an exception, or (None, None, None) otherwise.
    """
    if code is None:
        return
    if mode in {"exec", "single"}:
        func = exec
    else:
        func = eval
    try:
        func(code, glb, loc)
        return (None, None, None)
    except BaseException:
        type, value, traceback = sys.exc_info()
        # strip off the current frame as the traceback should only show user code
        traceback = traceback.tb_next
        return type, value, traceback


def get_cache_filename(fname, code=True):
    """
    Return the filename of the cache for the given filename.

    Cache filenames are similar to those used by the Mercurial DVCS for its
    internal store.

    The ``code`` switch should be true if we should use the code store rather
    than the script store.
    """
    datadir = XSH.env["XONSH_DATA_DIR"]
    cachedir = os.path.join(
        datadir, "xonsh_code_cache" if code else "xonsh_script_cache"
    )
    cachefname = os.path.join(cachedir, *_cache_renamer(fname, code=code))
    return cachefname


def update_cache(ccode, cache_file_name):
    """
    Update the cache at ``cache_file_name`` to contain the compiled code
    represented by ``ccode``.
    """
    if cache_file_name is not None:
        try:
            os.makedirs(os.path.dirname(cache_file_name), exist_ok=True)
        except OSError:
            writable = False
        else:
            writable = is_writable_file(cache_file_name)
        if not writable:
            if XSH.env.get("XONSH_DEBUG", "False"):
                print_warning(
                    f"update_cache: Cache file is not writable: {cache_file_name}\n"
                    f"Set $XONSH_CACHE_SCRIPTS=0, $XONSH_CACHE_EVERYTHING=0 to disable cache."
                )
            return
        with open(cache_file_name, "wb") as cfile:
            cfile.write(XONSH_VERSION.encode() + b"\n")
            cfile.write(bytes(PYTHON_VERSION_INFO_BYTES) + b"\n")
            marshal.dump(ccode, cfile)


def _check_cache_versions(cfile):
    # version data should be < 1 kb
    ver = cfile.readline(1024).strip()
    if ver != XONSH_VERSION.encode():
        return False
    ver = cfile.readline(1024).strip()
    return ver == PYTHON_VERSION_INFO_BYTES


def compile_code(filename, code, execer, glb, loc, mode):
    """
    Wrapper for ``execer.compile`` to compile the given code
    """

    if filename.endswith(".py") and mode == "exec":
        return compile(code, filename, mode)

    if not code.endswith("\n"):
        code += "\n"
    old_filename = execer.filename
    try:
        execer.filename = filename
        ccode = execer.compile(code, glbs=glb, locs=loc, mode=mode, filename=filename)
    except Exception:
        raise
    finally:
        execer.filename = old_filename
    return ccode


def script_cache_check(filename, cachefname):
    """
    Check whether the script cache for a particular file is valid.

    Returns a tuple containing: a boolean representing whether the cached code
    should be used, and the cached code (or ``None`` if the cache should not be
    used).
    """
    ccode = None
    run_cached = False
    if os.path.isfile(cachefname):
        if os.stat(cachefname).st_mtime >= os.stat(filename).st_mtime:
            with open(cachefname, "rb") as cfile:
                if not _check_cache_versions(cfile):
                    return False, None
                try:
                    ccode = marshal.load(cfile)
                except Exception:
                    # Cache file is corrupted (e.g. truncated by a crash).
                    # Ignore it — the script will be recompiled and cached again.
                    return False, None
                run_cached = True
    return run_cached, ccode


def run_script_with_cache(filename, execer, glb=None, loc=None, mode="exec"):
    """
    Run a script, using a cached version if it exists (and the source has not
    changed), and updating the cache as necessary.
    See run_compiled_code for the return value.
    """
    run_cached = False
    use_cache = should_use_cache(execer, mode)
    cachefname = get_cache_filename(filename, code=False)
    if use_cache:
        run_cached, ccode = script_cache_check(filename, cachefname)
    if not run_cached:
        with open(filename, encoding="utf-8") as f:
            code = f.read()
        ccode = compile_code(filename, code, execer, glb, loc, mode)
        if use_cache:
            update_cache(ccode, cachefname)
    return run_compiled_code(ccode, glb, loc, mode)


def code_cache_name(code):
    """
    Return an appropriate spoofed filename for the given code.
    """
    if isinstance(code, str):
        code = code.encode()
    # usedforsecurity=False: allow md5 on FIPS-enabled systems
    return hashlib.md5(code, usedforsecurity=False).hexdigest()


def code_cache_check(cachefname):
    """
    Check whether the code cache for a particular piece of code is valid.

    Returns a tuple containing: a boolean representing whether the cached code
    should be used, and the cached code (or ``None`` if the cache should not be
    used).
    """
    ccode = None
    run_cached = False
    if os.path.isfile(cachefname):
        with open(cachefname, "rb") as cfile:
            if not _check_cache_versions(cfile):
                return False, None
            try:
                ccode = marshal.load(cfile)
            except Exception:
                # Cache file is corrupted (e.g. truncated by a crash).
                # Ignore it — the code will be recompiled and cached again.
                return False, None
            run_cached = True
    return run_cached, ccode


def run_code_with_cache(
    code, display_filename, execer, glb=None, loc=None, mode="exec"
):
    """
    Run a piece of code, using a cached version if it exists, and updating the
    cache as necessary.
    See run_compiled_code for the return value.
    """
    use_cache = should_use_cache(execer, mode)
    filename = code_cache_name(code)
    cachefname = get_cache_filename(filename, code=True)
    run_cached = False
    if use_cache:
        run_cached, ccode = code_cache_check(cachefname)
    if not run_cached:
        ccode = compile_code(display_filename, code, execer, glb, loc, mode)
        if use_cache:
            update_cache(ccode, cachefname)
    return run_compiled_code(ccode, glb, loc, mode)
