#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import argparse
import logging
import os
import shutil
import site
import sys
from pathlib import Path


logging.basicConfig(level=logging.INFO)
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("project_name")
    parser.add_argument("--template-name", default="synthetic")
    args = parser.parse_args()

    root = Path(site.getsitepackages()[0])
    egg_link = root / "classy-vision.egg-link"
    dev_install = False
    # Support development mode (pip install -e)
    if egg_link.exists():
        dev_install = True
        with egg_link.open("r") as f:
            lines = f.read().split("\n")
            if lines[1] != ".":
                raise RuntimeError("Unexpected egg-link format")
            root = Path(lines[0])

    base_path = root / "classy_vision"
    template_path = base_path / "templates" / args.template_name
    destination_path = Path(os.getcwd()) / args.project_name

    if destination_path.exists():
        logging.error(f"Project directory '{destination_path}' already exists!")
        sys.exit(1)

    if dev_install:
        classy_train_path = root / "classy_train.py"
    else:
        classy_train_path = Path(sys.prefix) / "classy_vision" / "classy_train.py"

    shutil.copytree(template_path, destination_path)
    shutil.copy(classy_train_path, destination_path)

    logging.info(
        f"""
    Successfully generated template project at '{destination_path}'.
    To get started, run:
        $ cd {args.project_name}
        $ ./classy_train.py --config configs/template_config.json"""
    )
