#!/bin/bash
set -euo pipefail

export HOST_UID="$(id -u)"
export HOST_GID="$(id -g)"
export HOST_USER="$(whoami)"
export PDK_ROOT=/cad/share/pdk

IMAGE_TAG="srampr/ee5311_iitm:jul_nov_2026"
SAVED_TAG="srampr/ee5311_iitm:saved"

ACTION="${1:-up}"
if [ "$#" -gt 0 ]; then
    shift
fi

mkdir -p "$HOME/ee5311"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "${SCRIPT_DIR}"

compose_yaml() {
cat <<YAML
name: ee5311
services:
  image:
    image: ${1}
    entrypoint: ["bash", "-c", "rm -f /tmp/.X*-lock /tmp/.X11-unix/X*; exec /entrypoint.sh"]
    environment:
      - HOST_UID=${HOST_UID}
      - HOST_GID=${HOST_GID}
      - HOST_USER=${HOST_USER}
    volumes:
      - ${HOME}/ee5311:/home/${HOST_USER}/ee5311
    ports:
      - "${VNC_HOST_PORT:-5999}:5901"
YAML
}

resolve_ref() {
    if docker image inspect "${SAVED_TAG}" >/dev/null 2>&1; then
        echo "${SAVED_TAG}"
    else
        echo "${IMAGE_TAG}"
    fi
}

get_cid() {
    compose_yaml "$(resolve_ref)" | docker compose -f - ps -q image
}

do_save() {
    local cid
    cid="$(get_cid)"
    if [ -z "${cid}" ]; then
        return 1
    fi
    docker unpause "${cid}" >/dev/null 2>&1 || true
    docker exec "${cid}" bash -c \
        "touch /home/${HOST_USER}/.ee5311_saved_state && chown ${HOST_UID}:${HOST_GID} /home/${HOST_USER}/.ee5311_saved_state"
    docker commit "${cid}" "${SAVED_TAG}" >/dev/null
    echo "Saved container ${cid:0:12} as ${SAVED_TAG}; '$0 up' will resume from it." >&2
}

case "$ACTION" in
    up)
        REF="$(resolve_ref)"
        if [ "${REF}" = "${IMAGE_TAG}" ]; then
            echo "No ${SAVED_TAG} snapshot yet; starting from ${IMAGE_TAG}." >&2
        fi
        compose_yaml "${REF}" | docker compose -f - up -d "$@"
        ;;
    fresh)
        compose_yaml "${IMAGE_TAG}" | docker compose -f - up -d --force-recreate "$@"
        ;;
    save)
        do_save || {
            echo "No running ee5311 container to save; start one with '$0 up' or '$0 fresh'." >&2
            exit 1
        }
        ;;
    down)
        if [ "${1:-}" = "--no-save" ]; then
            shift
            CID="$(get_cid)"
            { [ -n "${CID}" ] && docker unpause "${CID}" >/dev/null 2>&1; } || true
        else
            do_save || echo "No running ee5311 container; nothing to save." >&2
        fi
        compose_yaml "$(resolve_ref)" | docker compose -f - down "$@"
        ;;
    pause)
        compose_yaml "$(resolve_ref)" | docker compose -f - pause image
        ;;
    resume)
        compose_yaml "$(resolve_ref)" | docker compose -f - unpause image
        ;;
    *)
        echo "Usage: $0 [up|fresh|save|down|pause|resume] [extra docker compose args...]" >&2
        echo "  up      start from the ${SAVED_TAG} snapshot if present, else ${IMAGE_TAG}" >&2
        echo "  fresh   force a new container from ${IMAGE_TAG} (ignores saved state)" >&2
        echo "  save    commit the running container's state to ${SAVED_TAG}" >&2
        echo "  down    save the running container (unless --no-save), then stop and remove the stack" >&2
        echo "  pause   freeze the running container's processes in RAM (resume to continue)" >&2
        echo "  resume  unfreeze a paused container" >&2
        exit 1
        ;;
esac