RESTART_FILE = '/tmp/.restart-proc'
TYPE_RESTART_CONTAINER_STEP = 'live_update_restart_container_step'

KWARGS_BLACKLIST = [
    # since we'll be passing `dockerfile_contents` when building the
    # child image, remove any kwargs that might conflict
    'dockerfile', 'dockerfile_contents',

    # 'target' isn't relevant to our child build--if we pass this arg,
    # Docker will just fail to find the specified stage and error out
    'target',
]

# Arguments to custom_build that don't apply to the docker_build.
_CUSTOM_BUILD_KWARGS_BLACKLIST = [
    'tag',
    'command_bat',
    'outputs_image_ref_to',
    'disable_push',
]

_ext_dir = os.getcwd()

# shared code between the two restart functions
def _helper(base_ref, ref, entrypoint, live_update, restart_file=RESTART_FILE, trigger=None, exit_policy='restart', **kwargs):
    if not trigger:
        trigger = []

    # declare a new docker build that adds a static binary of tilt-restart-wrapper
    # (which makes use of `entr` to watch files and restart processes) to the user's image
    df = '''
    FROM tiltdev/restart-helper:2021-08-09 as restart-helper

    FROM {}
    RUN ["touch", "{}"]
    COPY --from=restart-helper /tilt-restart-wrapper /
    COPY --from=restart-helper /entr /
    '''.format(base_ref, restart_file)

    # Change the entrypoint to use `tilt-restart-wrapper`.
    # `tilt-restart-wrapper` makes use of `entr` (https://github.com/eradman/entr/) to
    # re-execute $entrypoint whenever $restart_file changes
    entrypoint_with_entr = ["/tilt-restart-wrapper", "--watch_file={}".format(restart_file)]
    if exit_policy == 'continue':
        entrypoint_with_entr = entrypoint_with_entr + ["--entr_flags=-r"]
    if type(entrypoint) == type(""):
        entrypoint_with_entr = entrypoint_with_entr + ["sh", "-c", entrypoint]
    elif type(entrypoint) == type([]):
        entrypoint_with_entr = entrypoint_with_entr + entrypoint
    else:
        fail("`entrypoint` must be a string or list of strings: got {}".format(type(entrypoint)))

    # last live_update step should always be to modify $restart_file, which
    # triggers the process wrapper to rerun $entrypoint
    # NB: write `date` instead of just `touch`ing because `entr` doesn't respond
    # to timestamp changes, only writes (see https://github.com/eradman/entr/issues/32)
    live_update = live_update + [run('date > {}'.format(restart_file), trigger=trigger)]

    # We don't need a real context. See:
    # https://github.com/tilt-dev/tilt/issues/3897
    context = _ext_dir

    docker_build(ref, context, entrypoint=entrypoint_with_entr, dockerfile_contents=df,
                 live_update=live_update, **kwargs)

def docker_build_with_restart(ref, context, entrypoint, live_update,
                              base_suffix='-tilt_docker_build_with_restart_base', restart_file=RESTART_FILE,
                              trigger=None, exit_policy='restart', **kwargs):
    """Wrap a docker_build call and its associated live_update steps so that the last step
    of any live update is to rerun the given entrypoint.

     Args:
      ref: name for this image (e.g. 'myproj/backend' or 'myregistry/myproj/backend'); as the parameter of the same name in docker_build
      context: path to use as the Docker build context; as the parameter of the same name in docker_build
      entrypoint: the command to be (re-)executed when the container starts or when a live_update is run
      live_update: set of steps for updating a running container; as the parameter of the same name in docker_build
      base_suffix: suffix for naming the base image, applied as {ref}{base_suffix}
      restart_file: file that Tilt will update during a live_update to signal the entrypoint to rerun
      trigger: (optional) list of local paths. If specified, the process will ONLY be restarted when there are changes
               to the given file(s); as the parameter of the same name in the LiveUpdate `run` step.
      **kwargs: will be passed to the underlying `docker_build` call
    """

    # first, validate the given live_update steps
    if len(live_update) == 0:
        fail("`docker_build_with_restart` requires at least one live_update step")
    for step in live_update:
        if type(step) == TYPE_RESTART_CONTAINER_STEP:
            fail("`docker_build_with_restart` is not compatible with live_update step: " +
                 "`restart_container()` (this extension is meant to REPLACE restart_container() )")

    # rename the original image to make it a base image and declare a docker_build for it
    base_ref = '{}{}'.format(ref, base_suffix)
    docker_build(base_ref, context, **kwargs)

    # Clean kwargs for building the child image (which builds on user's specified
    # image and copies in Tilt's restart wrapper). In practice, this means removing
    # kwargs that were relevant to building the user's specified image but are NOT
    # relevant to building the child image / may conflict with args we specifically
    # pass for the child image.
    cleaned_kwargs = {k: v for k, v in kwargs.items() if k not in KWARGS_BLACKLIST}
    _helper(base_ref, ref, entrypoint, live_update, restart_file, trigger, exit_policy, **cleaned_kwargs)


def custom_build_with_restart(ref, command, deps, entrypoint, live_update,
                              base_suffix='-tilt_docker_build_with_restart_base', restart_file=RESTART_FILE,
                              trigger=None, exit_policy='restart', **kwargs):
    """Wrap a custom_build call and its associated live_update steps so that the last step
    of any live update is to rerun the given entrypoint.

     Args:
      ref: name for this image (e.g. 'myproj/backend' or 'myregistry/myproj/backend'); as the parameter of the same name in custom_build
      command: build command for building your image
      deps: source dependencies of the custom build
      entrypoint: the command to be (re-)executed when the container starts or when a live_update is run
      live_update: set of steps for updating a running container; as the parameter of the same name in custom_build
      base_suffix: suffix for naming the base image, applied as {ref}{base_suffix}
      restart_file: file that Tilt will update during a live_update to signal the entrypoint to rerun
      trigger: (optional) list of local paths. If specified, the process will ONLY be restarted when there are changes
               to the given file(s); as the parameter of the same name in the LiveUpdate `run` step.
      **kwargs: will be passed to the underlying `custom_build` call
    """

    # first, validate the given live_update steps
    if len(live_update) == 0:
        fail("`custom_build_with_restart` requires at least one live_update step")
    for step in live_update:
        if type(step) == TYPE_RESTART_CONTAINER_STEP:
            fail("`custom_build_with_restart` is not compatible with live_update step: "+
                 "`restart_container()` (this extension is meant to REPLACE restart_container() )")

    for k, v in kwargs.items():
        if k == 'skips_local_docker':
            fail("`custom_build_with_restart` is not compatible with `skips_local_docker`, because it needs access to the image")
        if k == 'disable_push':
            fail("`custom_build_with_restart` is not compatible with `disable_push`")
        if k == 'tag':
            fail("`custom_build_with_restart` renames your base image, so is not compatible with `tag`")

    # rename the original image to make it a base image and declare a custom_build for it
    base_ref = '{}{}'.format(ref, base_suffix)
    custom_build(base_ref, command=command, deps=deps, **kwargs)

    # A few arguments aren't applicable to the docker_build, so remove them.
    cleaned_kwargs = {k: v for k, v in kwargs.items() if k not in _CUSTOM_BUILD_KWARGS_BLACKLIST}
    _helper(base_ref, ref, entrypoint, live_update, restart_file, trigger, exit_policy, **cleaned_kwargs)
