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'

#TODO: load from config file
DEFAULT_LANG = 'en'
OTHER_LANGS = set(['es'])


@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")


def add_lang_prefix(lang, path):
    if lang == DEFAULT_LANG:
        return path

    if lang not in OTHER_LANGS:
        print(f"Not registered language: `{lang}`")

    return f"/{lang}{path}"


@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, lang):
    name, _ = splitext(ctx.name)
    if name.startswith('./'):
        name = name[2:]

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

    return add_lang_prefix(lang, f"/{name}")


@contextfilter
def cur_lang(ctx, path):
    lang = ctx.environment.globals['lang']
    return add_lang_prefix(lang, path)


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['cur_lang'] = cur_lang
    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 OTHER_LANGS:
        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)