# based on https://github.com/hgomersall/scons-jinja

from SCons.Script import *

import os
import pathlib

from spp import get_spp

try:
    import jinja2
    from jinja2.utils import open_if_exists
except ImportError:
    jinja2 = None
    print('No Jinja :(')

spp = get_spp()

def available(**kwargs) -> bool:
    return jinja2 is not None

def post_environment(**kwargs) -> None:
    env: Environment = spp.globals['env']

    env.SetDefault(JINJA_CONTEXT={})
    env.SetDefault(JINJA_ENVIRONMENT_VARS={})
    env.SetDefault(JINJA_FILTERS={'load_config': _jinja_load_config})
    env.SetDefault(JINJA_GLOBALS={
        'file_size': lambda *args: _file_size(env, *args),
        'file_content_hex': lambda *args: _file_content_hex(env, *args)
    })
    env.SetDefault(JINJA_TEMPLATE_SEARCHPATH=['data/jinja'])
    env.SetDefault(JINJA_CONFIG_SEARCHPATH=[env.Dir('#data/config')])
    env.SetDefault(JINJA_FILE_SEARCHPATH=[env.Dir('#')])

    env['BUILDERS']['Jinja'] = Builder(
        action=render_jinja_template
    )

    scanner = env.Scanner(function=jinja_scanner,
                          skeys=['.jinja'])

    env.Append(SCANNERS=scanner)

    env.AddMethod(_wrap_jinja(env.Jinja), 'Jinja')

class FileSystemLoaderRecorder(jinja2.FileSystemLoader):
    """ A wrapper around FileSystemLoader that records files as they are
    loaded. These are contained within loaded_filenames set attribute.
    """

    def __init__(self, searchpath, encoding='utf-8'):

        self.loaded_filenames = set()
        super(FileSystemLoaderRecorder, self).__init__(searchpath, encoding)

    def get_source(self, environment, template):
        """Overwritten FileSystemLoader.get_source method that extracts the
        filename that is used to load each filename and adds it to
        self.loaded_filenames.
        """
        for searchpath in self.searchpath:
            filename = os.path.join(searchpath, template)
            f = open_if_exists(filename)
            if f is None:
                continue
            try:
                contents = f.read().decode(self.encoding)
            finally:
                f.close()

            self.loaded_filenames.add(filename)

            return super(FileSystemLoaderRecorder, self).get_source(
                environment, template)

        # If the template isn't found, then we have to drop out.
        raise jinja2.TemplateNotFound(template)


def jinja_scanner(node, env, path):
    # Instantiate the file as necessary
    node.get_text_contents()

    template_dir, filename = os.path.split(str(node))

    template_search_path = ([template_dir] +
                            env.subst(env['JINJA_TEMPLATE_SEARCHPATH']))
    template_loader = FileSystemLoaderRecorder(template_search_path)

    jinja_env = jinja2.Environment(loader=template_loader,
                                   extensions=['jinja2.ext.do'], **env['JINJA_ENVIRONMENT_VARS'])
    jinja_env.filters.update(env['JINJA_FILTERS'])
    jinja_env.globals.update(env['JINJA_GLOBALS'])

    try:
        template = jinja_env.get_template(filename)
    except jinja2.TemplateNotFound as e:
        env.Error(f'Missing template: {os.path.join(template_dir, str(e))}')

    # We need to render the template to do all the necessary loading.
    #
    # It's necessary to respond to missing templates by grabbing
    # the content as the exception is raised. This makes sure of the
    # existence of the file upon which the current scanned node depends.
    #
    # I suspect that this is pretty inefficient, but it does
    # work reliably.
    context = env['JINJA_CONTEXT']

    last_missing_file = ''
    while True:

        try:
            template.render(**context)
        except jinja2.TemplateNotFound as e:
            if last_missing_file == str(e):
                # We've already been round once for this file,
                # so need to raise
                env.Error(f'Missing template: {os.path.join(template_dir, str(e))}')

            last_missing_file = str(e)
            # Find where the template came from (using the same ordering
            # as Jinja uses).
            for searchpath in template_search_path:
                filename = os.path.join(searchpath, last_missing_file)
                if os.path.exists(filename):
                    continue
                else:
                    env.File(filename).get_text_contents()
            continue

        break

    # Get all the files that were loaded. The set includes the current node,
    # so we remove that.
    found_nodes_names = list(template_loader.loaded_filenames)
    try:
        found_nodes_names.remove(str(node))
    except ValueError as e:
        env.Error(f'Missing template node: {str(node)}')

    return [env.File(f) for f in found_nodes_names]


def render_jinja_template(target, source, env):
    output_str = ''

    if not source:
        source = [f'{target}.jinja']

    for template_file in source:
        template_dir, filename = os.path.split(str(template_file))

        template_search_path = ([template_dir] +
                                env.subst(env['JINJA_TEMPLATE_SEARCHPATH']))
        template_loader = FileSystemLoaderRecorder(template_search_path)

        jinja_env = jinja2.Environment(loader=template_loader,
                                       extensions=['jinja2.ext.do'], **env['JINJA_ENVIRONMENT_VARS'])
        jinja_env.filters.update(env['JINJA_FILTERS'])
        jinja_env.globals.update(env['JINJA_GLOBALS'])

        jinja_env.filters.update(env['JINJA_FILTERS'])
        template = jinja_env.get_template(filename)

        context = env['JINJA_CONTEXT']
        template.render(**context)

        output_str += template.render(**context)

    with open(str(target[0]), 'w') as target_file:
        target_file.write(output_str)

    return None

def _jinja_load_config(env, config_name):
    searched_paths = []
    for scons_path in env['JINJA_CONFIG_SEARCHPATH']:
        if hasattr(scons_path, 'abspath'):
            scons_path = scons_path.abspath
        path = pathlib.Path(scons_path) / f'{config_name}.yml'
        if path.exists():
            with path.open('r') as file:
                import yaml
                return yaml.safe_load(file)
        searched_paths.append(f'\n{path}')
    joined_paths = ''.join(searched_paths)
    raise Exception(f'Could not find Jinja config file "{config_name}.yml". Searched: {joined_paths}')
def _wrap_jinja(orig_jinja):
    def _wrapped(env, target, **kwargs):
        if 'source' not in kwargs:
            kwargs['source'] = f'{target}.jinja'
        target = orig_jinja(target=target, **kwargs)
        if 'depends' in kwargs:
            for dependency in kwargs['depends']:
                env.Depends(target, dependency)
        return target
    return _wrapped

def _find_file(env, fname):
    for path in env['JINJA_FILE_SEARCHPATH']:
        fullpath = os.path.join(path.abspath, fname)
        if os.path.exists(fullpath):
            return env.File(fullpath)
    return None

def _file_size(env, fname: str) -> int:
    file = _find_file(env, fname)
    if not file:
        env.Error(f'File does not exist: {fname}. Searched in: {[d.abspath for d in env["JINJA_FILE_SEARCHPATH"]]}')
    return file.get_size()

def _file_content_hex(env, fname: str) -> str:
    file = _find_file(env, fname)
    if not file:
        env.Error(f'File does not exist: {fname}. Searched in: {[d.abspath for d in env["JINJA_FILE_SEARCHPATH"]]}')
    bytes = file.get_contents()
    return ','.join([hex(byte) for byte in bytes])