#!/usr/bin/python3

import os
from string import Template


DEFAULT_FIREJAIL_PATH = os.path.join(os.path.expanduser('~'), '.config', 'firejail')
DEFAULT_TEMPLATE_PATH = os.path.join(
    os.path.dirname(os.path.realpath(__file__)),
    'python-env-template.profile',
)
DEFAULT_ENVS_PATH = os.path.join(os.path.expanduser('~'), 'envs')


class ProfileTemplate(Template):
    delimiter = '€'


def read_template(path):
    with open(path) as f:
        return ProfileTemplate(f.read())


def write_profile(path, content):
    with open(path, 'w') as f:
        f.write(content)


def write_sandbox_dotfile(sandbox_name, working_dir):
    """This allows for shells to automatically enable the
    sandbox once they find the file.
    """
    path = os.path.join(working_dir, '.python-sandbox')
    with open(path, 'w') as f:
        f.write(sandbox_name + '\n')
    return path


def resolve_working_dir(working_dir):
    """Make sure to get an absolute path or raise an error.
    For example, if "." is supplied, it should resolve in the
    current working dir as an absolute path since `.` would
    not be valid as working dir in the template.
    """
    return os.path.abspath(working_dir)


def resolve_env_dir():
    """Make sure that the user can permanently overwrite the place
    where python environments are assumed to reside by the firejail
    profile.

    Method: Set `PYSANDBOX_ENVS_PATH` environment variable.
    """
    return os.environ.get('PYSANDBOX_ENVS_PATH', DEFAULT_ENVS_PATH)


def main(args):
    print(f'Loading template {args.template}.')

    working_dir = resolve_working_dir(args.working_dir)

    template_str = read_template(args.template)
    profile_str = template_str.substitute({
        'sandbox_name': args.sandbox_name,
        'working_dir': working_dir,
        'env_dir': resolve_env_dir(),
        'HOME': os.path.expanduser('~'),
    })

    if not os.path.exists(DEFAULT_FIREJAIL_PATH):
        os.makedirs(DEFAULT_FIREJAIL_PATH, exist_ok=True)

    profile_name = f'python-env-{args.sandbox_name}.profile'
    profile_path = os.path.join(DEFAULT_FIREJAIL_PATH, profile_name)

    if os.path.exists(profile_path) and not args.force:
        print('Profile already exists. Aborting.')
        return 1

    write_profile(profile_path, profile_str)

    sandbox_dotfile = write_sandbox_dotfile(
        args.sandbox_name,
        working_dir,
    )

    print(f"Written profile to {profile_path}.")
    print(f"Written {args.sandbox_name} to {sandbox_dotfile}.")
    print("")
    print(f"""Start your sandbox by running

    firejail --profile=python-env-{args.sandbox_name} --tab bash

""")


if __name__ == "__main__":
    import sys
    from argparse import ArgumentParser
    parser = ArgumentParser(
        description="""Create a new python sandbox profile in the firejail
configuration directory given the name of the sandbox
and it's working directory (which gets full read-write)
access. The tool automatically places a `.python-sandbox` file with the
sandbox's name in the working directory.

For venv/conda/... environments, by default ~/envs/<sandbox name> is allowed.
To override the default directory where environments are supposed to be
stored, set the PYSANDBOX_ENVS_PATH environment variable.
"""
    )
    parser.add_argument('sandbox_name', type=str, help="name of sandbox (needs to be FS compatible)")
    parser.add_argument('working_dir', type=str, help="Working directory of the sandbox")
    parser.add_argument('--template', type=str, default=DEFAULT_TEMPLATE_PATH, help="Path to the firejail profile template")
    parser.add_argument('-f', '--force', action='store_true', help="Force overriding the sandbox if it already exists")

    args = parser.parse_args()
    ret = main(args)
    sys.exit(ret or 0)