import os
from typing import Union
import yaml
import json
from jinja2 import Template
from copy import deepcopy

header_msg1 = 'This file was automatically generated by jinjaroot. Do not edit directly. See the .jinjaroot dir.'
header_msg2 = 'This file was automatically generated by jinjaroot. Do not edit directly.'

def process_jinjaroot_dir(*, src_dirpath: str, dst_dirpath: str, template_data: dict, verify_only: bool, dry: bool):
    template_fnames = os.listdir(src_dirpath)
    for fname in template_fnames:
        template_path = f'{src_dirpath}/{fname}'
        fname2 = Template(fname).render(**template_data)
        dest_path = f'{dst_dirpath}/{fname2}'
        if os.path.isdir(template_path):
            if not os.path.exists(dest_path):
                os.mkdir(dest_path)
            process_jinjaroot_dir(src_dirpath=template_path, dst_dirpath=dest_path, template_data=template_data, verify_only=verify_only, dry=dry)
        elif os.path.isfile(template_path):
            with open(template_path, 'r') as f:
                template_code = f.read()
            t = Template(template_code, keep_trailing_newline=True)
            gen_code = t.render(**template_data)
            gen_code = add_gen_header(fname, gen_code, header_msg1)
            _check_and_write_file(dst_path=dest_path, gen_code=gen_code, verify_only=verify_only, dry=dry, template_path=template_path)

def _check_and_write_file(*, dst_path: str, gen_code: str, verify_only: bool, dry: bool, template_path: Union[str, None]):
    if template_path is not None:
        template_stat = os.stat(template_path)
    else:
        template_stat = None
    if os.path.exists(dst_path):
        with open(dst_path, 'r') as f:
            dst_code = f.read()
        dst_stat = os.stat(dst_path)
    else:
        dst_code = None
        dst_stat = None
    if gen_code != dst_code:
        if verify_only:
            raise Exception(f'Code generation is not up-to-date for file: {dst_path}')
        if dry:
            print(f'Will update: {dst_path}')
        else:
            print(f'Writing: {dst_path}')
            with open(dst_path, 'w') as f:
                f.write(gen_code)
            if (dst_stat is None) or (template_stat and (template_stat.st_mode != dst_stat.st_mode)):
                os.chmod(dst_path, template_stat.st_mode)

def process_dir(*, dirpath: str, template_data: dict, verify_only: bool, dry: bool):
    jinja_root_yaml_path = f'{dirpath}/jinjaroot.yaml'
    jinja_root_json_path = f'{dirpath}/jinjaroot.json'
    jinjaroot_dirpath = f'{dirpath}/.jinjaroot'
    template_data2 = deepcopy(template_data)
    if os.path.isfile(jinja_root_yaml_path):
        with open(jinja_root_yaml_path, 'r') as f:
            x = yaml.safe_load(f)
            for k, v in x.items():
                template_data2[k] = v
    if os.path.isfile(jinja_root_json_path):
        with open(jinja_root_json_path, 'r') as f:
            x = json.load(f)
            for k, v in x.items():
                template_data2[k] = v
    if os.path.isdir(jinjaroot_dirpath):
        process_jinjaroot_dir(src_dirpath=jinjaroot_dirpath, dst_dirpath=dirpath, template_data=template_data2, verify_only=verify_only, dry=dry)
    
    fnames = os.listdir(dirpath)
    for fname in fnames:
        file_path = f'{dirpath}/{fname}'
        if os.path.isfile(file_path):
            template_path = file_path + '.j2'
            if os.path.exists(template_path):
                with open(template_path, 'r') as f:
                    template_code = f.read()
                t = Template(template_code, keep_trailing_newline=True)
                gen_code = t.render(**template_data2)
                gen_code = add_gen_header(fname, gen_code, header_msg2)
                _check_and_write_file(dst_path=file_path, gen_code=gen_code, verify_only=verify_only, dry=dry, template_path=template_path)
            if ends_with(file_path, ['.json']):
                substitute_template_path = file_path + '.substitute.j2'
                if os.path.exists(substitute_template_path):
                    with open(file_path, 'r') as f:
                        content1 = json.load(f)
                    with open(substitute_template_path, 'r') as f:
                        txt = f.read()
                        txt2 = Template(txt, keep_trailing_newline=True).render(**template_data2)
                        content2 = json.loads(txt2)
                    content1_new = _substitute_json(content1, content2)
                    gen_code = json.dumps(content1_new, indent=2)
                    _check_and_write_file(dst_path=file_path, gen_code=gen_code, verify_only=verify_only, dry=dry, template_path=None)
        elif os.path.isdir(file_path):
            process_dir(dirpath=file_path, template_data=template_data2, verify_only=verify_only, dry=dry)

def _substitute_json(content1: dict, content2: dict):
    ret = deepcopy(content1)
    for k, v in content2.items():
        ret[k] = v
    return ret

def ends_with(fname, suffixes):
    for suf in suffixes:
        if fname.endswith(suf): return True
    return False

def insert_header(header: str, code: str):
    lines = code.splitlines()
    line1 = lines[0] if len(lines) > 0 else ''
    if line1.startswith('#!'):
        return code.replace(line1, line1 + '\n' + header)
    else:
        return header + code

def add_gen_header(fname: str, code: str, message: str):
    if ends_with(fname, ['.js', '.jsx', '.ts', '.tsx']):
        return insert_header(f'// {message}\n', code)
    elif ends_with(fname, ['.py', '.sh', '.yaml', '.yml', '.cfg']):
        return insert_header(f'# {message}\n', code)
    elif ends_with(fname, ['.html']):
        return insert_header(f'<!-- {message} -->\n', code)
    elif ends_with(fname, ['.md']):
        return insert_header(f'[//]: # "{message}"\n', code)
    else:
        return code


def generate(*, verify_only: bool=False, dry: bool=False):
    template_data = {}
    process_dir(dirpath='.', template_data=template_data, verify_only=verify_only, dry=dry)

def verify():
    generate(verify_only=True, dry=False)
