#!/usr/bin/env python3

import sys
import os
import re
import shutil
import argparse
import tempfile
import subprocess
from glob import glob
from importlib import import_module
from jinja2 import Template, Environment, PackageLoader, FileSystemLoader

from modelmaker import folders_in_folder, make_directory

module_folder = os.path.dirname(sys.modules['modelmaker'].__file__)
current_folder = os.path.abspath(os.curdir)

templates = { os.path.basename(t): t for t in folders_in_folder(os.path.join(module_folder, 'resources', 'templates')) }

# parsers
parser = argparse.ArgumentParser(description='modelmaker tool for creating and validating ML models')
subparsers = parser.add_subparsers(dest='subcommand')

# parsers - new
new_project_parser = subparsers.add_parser('new', description='create a new project in current directory')
new_project_parser.add_argument('--project', help='project and main class name')
new_project_parser.add_argument('--package', help='python pacake name')
new_project_parser.add_argument('--template', default='default', help='selected template')

# parsers - templates
templates_parser = subparsers.add_parser('templates', description='list existing templates')
args = parser.parse_args()

############################################################################################
## create a new project
############################################################################################

if args.subcommand == 'new':

    project_path = os.path.join(current_folder, args.project)

    if os.path.isdir(project_path):
        print(f"{project_path} already exists")
        sys.exit(1)

    template = args.template
    template_folder = os.path.join(module_folder, 'resources', 'templates', template)

    if not os.path.isdir(template_folder):
        print(f"{template} template does not exist ...")
        sys.exit(1)

    print(f"Creating {template} template in {project_path} ...")

    # copy project template
    shutil.copytree(
        template_folder,
        project_path,
        ignore=shutil.ignore_patterns('*.pyc', '__pycache__')
    )

    # rename package folder
    os.rename(
        os.path.join(project_path, 'package'),
        os.path.join(project_path, args.package)
    )

    # create .gitignore
    with open(os.path.join(project_path, '.gitignore'), 'w') as f:
        f.write('saved_models')

    # replace template strings
    env = Environment(loader=FileSystemLoader(project_path))
    replace = {
        'package_name': args.package,
        'project_name': args.project
    }
    for file in glob(os.path.join(project_path, '**/*'), recursive=True):
        if os.path.isfile(file):
            template = env.get_template(os.path.relpath(file, start=project_path))
            with open(file, 'w') as f:
                print(f"Generating template {file} ...")
                f.write(template.render(**replace))

############################################################################################
## list templates
############################################################################################

elif args.subcommand == 'templates':
    print("\nTemplates available:")
    for t in templates:
        print(f"    {t}")
    print("\n")

else:
    parser.print_help()
