import os
import jinja2
import argparse
import traceback

def render_template(template_file: str, 
                    out_path: str,
                    **kwargs) -> None:
    folder, file = os.path.split(template_file)
    template_path = os.path.abspath(folder)
    env = jinja2.Environment(loader=jinja2.FileSystemLoader(template_path),
                             trim_blocks=True,
                             lstrip_blocks=True)
    template = env.get_template(file)
    with open(out_path, 'w+') as save:
        save.write(template.render(**kwargs))


def load_template_vars_file(var_file_path: str) -> dict:
    result = {}
    with open(var_file_path) as reader:
        for line in reader:
            line = line.strip()
            if line.startswith('#') or '=' not in line:
                continue
            key, value = line.split('=', 1)
            result[key] = value

    return result


if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='render_templates.py',
                                     description='A script to render the docker jinja templates.',
                                     add_help=True)
    parser.add_argument('template_path',
                        help='The path to the template(s) to render (Can be either a directory or a file).')
    parser.add_argument('var_file_path',
                        help='The path to the file containing the variables to use for rendering the tempalte(s).')
    parser.add_argument('out_path',
                        help='The path to write the rendered template (MUST be a directory).')
    parser.add_argument('-k', 
                        '--kwargs',
                        required=False,
                        type=str,
                        default='',
                        help='Additional kwargs to include if needed (format: key1=value1,key2=value2...).')
    parser.add_argument('-o', 
                        '--overwrite',
                        action='store_true',
                        help='If set the previous template will be overwritten.')

    args = parser.parse_args()

    # make sure we have a template file and a variable file.
    if not os.path.exists(args.template_path):
        print(f'The given template path {args.template_path} does not exist.')
        exit(1)
    
    if not os.path.exists(args.var_file_path):
        print(f'The given template path {args.var_file_path} does not exist.')
        exit(1)
    
    run_on_dir = False
    # if the template path is a directory, make sure output path exists
    if os.path.isdir(args.template_path):
        run_on_dir = True
        if not os.path.exists(args.out_path):
            os.makedirs(args.out_path)
        else:
            print('The output path must be a directory since the template path is.')
            exit(1)

        args.template_path = [os.path.join(args.template_path, p) 
                              for p in filter(lambda x: x.ends_with('.tmpl'),
                                              os.listdir(args.template_path))]
    # if the tempalte path is a file
    else:
        # wrap it in a list
        args.template_path = [args.template_path]

    # load the variables from the variable file
    kwargs = load_template_vars_file(args.var_file_path)
    args.kwargs = {a.split('=')[0]: a.split('=')[1] for a in args.kwargs.split(',')} if args.kwargs else {}
    for k in kwargs:
        # The assumption made here is that anything passed on the commandline should be prioritized
        # over values that appear in the template variables file.
        if k not in args.kwargs:
            args.kwargs[k] = kwargs[k]

    for template in args.template_path:
        out_file_name = os.path.join(args.out_path,
                                     os.path.split(template.replace('.tmpl', ''))[1])
        # if the output exists and we aren't overwritting files, then move on.
        if os.path.exists(out_file_name) and not args.overwrite:
            continue

        try:
            render_template(template, out_file_name, **args.kwargs)
        except Exception as ex:
            print(f'Failed to render template "{template}":')
            print(f'Error: {str(ex)}\n{traceback.format_exec()}')