import argparse
from jinja2 import Environment, FileSystemLoader, contextfilter
from pathlib import Path
from os import walk, makedirs
from os.path import join, exists, splitext, split
from shutil import rmtree, copy, copytree
from functools import lru_cache

LANGUAGES = 'languages'

DEFAULT_LANG = 'en'
REGISTERED_LANG = set()


@lru_cache(None)
def get_lang_dic(lang):
    if not exists(LANGUAGES):
        makedirs(LANGUAGES)

    lang_file = join(LANGUAGES, lang + '.txt')

    if not exists(lang_file):
        return {}
    else:
        dic = {}
        with open(lang_file) as f:
            for line in f.readlines():
                key, value = line.strip('\n').split(':')
                dic[key] = value
        return dic


def save_dic(lang, dic):
    with open(join(LANGUAGES, lang + '.txt'), 'w') as f:
        for key, value in dic.items():
            f.write(f"{key}:{value}\n")


@contextfilter
def lang(ctx, value):
    lang = ctx.environment.globals['lang']

    if lang == DEFAULT_LANG:
        return value

    else:
        dic = get_lang_dic(lang)

        if value in dic and dic[value]:
            return dic[value]
        else:
            dic[value] = ""
            save_dic(lang, dic)
            print(f"Not translated phrase: `{value}`")
            return value


@contextfilter
def lang_url(ctx, value):
    name, _ = splitext(ctx.name)
    if name.startswith('./'):
        name = name[2:]

    if name == 'index':
        name = ''

    if value == DEFAULT_LANG:
        return "/" + name
    else:
        REGISTERED_LANG.add(value)
        return "/" + value + '/' + name


def static(value):
    return f"/static/{value}"


def compile(env, path, target):
    lang = env.globals['lang']

    for base, _, docs in walk(path):
        cur_base = Path(base)
        cur_base = Path(*cur_base.parts[1:])
        for doc in docs:
            if doc.startswith('_'):
                # Ignore this files
                continue

            src = join(base, doc)

            if args.with_index and doc != 'index.html':
                name, _ = splitext(doc)
                dst = join(cur_base, name, 'index.html')
            else:
                dst = join(cur_base, doc)

            if lang != DEFAULT_LANG:
                dst = join(lang, dst)

            dst = join(target, dst)
            dst_folder, _ = split(dst)

            if not exists(dst_folder):
                makedirs(dst_folder)

            template = env.get_template(join(cur_base, doc))
            output = template.render()

            with open(dst, 'w') as f:
                f.write(output)


def run(args):
    # Load layout
    file_loader = FileSystemLoader(args.source)
    env = Environment(loader=file_loader)
    env.globals['lang'] = DEFAULT_LANG

    # Add filters
    env.filters['lang'] = lang
    env.filters['lang_url'] = lang_url
    env.filters['static'] = static

    path = Path(args.source)

    if exists(args.target):
        rmtree(args.target)

    if exists('root'):
        copytree('root', args.target)
    else:
        makedirs(args.target)

    compile(env, path, args.target)

    copytree('static', join(args.target, 'static'))

    for other_lang in REGISTERED_LANG:
        env.globals['lang'] = other_lang
        compile(env, path, args.target)


if __name__ == '__main__':
    parser = argparse.ArgumentParser("Site generator")

    parser.add_argument('--with_index', action='store_true', default=False)
    parser.add_argument('--source', default='layout')
    parser.add_argument('--target', default='build')

    args = parser.parse_args()
    run(args)