#!/usr/bin/env bash

set -eu -o pipefail

SCRIPT_PATH=$(readlink -f "${BASH_SOURCE[0]}")
SCRIPT_NAME=$(basename "${BASH_SOURCE[0]}")

MOUNTABLE_SCRIPT_PATH="$HOME/.cache/$SCRIPT_NAME/$SCRIPT_NAME"
CONTAINER_SCRIPT_PATH="/sandbox/bin/$SCRIPT_NAME"
PRESET_BASE_IMAGE="public.ecr.aws/docker/library/debian:stable-20260505-slim"

help() {
  cat << HELP
$SCRIPT_NAME - Run a command in a sandboxed Docker environment

Usage: $SCRIPT_NAME [--dockerfile FILE]
           [--platform PLATFORM]
           [--env-file FILE]
           [--allow-write] [--allow-net [DESTINATIONS|--]]
           [--volume [NAME:]PATH]
           [--mount-writable HOST_DIR:CONTAINER_DIR]
           [--mount-readonly HOST_DIR:CONTAINER_DIR]
           [--publish [HOST_ADDRESS:]HOST_PORT:CONTAINER_PORT]
           [--tty]
           [--no-cache]
           [--skip-build]
           [--verbose] [--dry-run]
           [--keep-alive SECONDS]
           COMMAND


Options:

  --dockerfile FILE         Path to a Dockerfile. If not set, a preset Dockerfile
                            and network policy is used.
                            The container must include busybox, bash, iptables, ipset,
                            dnsmasq, and dig commands.
  --env-file FILE           Path to a .env file to load environment variables from.
  --platform                Specify the platform for image building and container execution
                            (e.g., linux/arm64 or linux/amd64).
  --allow-write             Allow write access to the project root inside the container.
                            By default, the project root (git root or current directory)
                            is read-only.
  --allow-net DESTINATIONS  Allow connections to specified domains or IP addresses.
                            Separate multiple destinations with commas.
                            If no port is given, only HTTPS (443) is allowed.
  --volume [NAME:]PATH      Mount a writable Docker volume at PATH inside the container.
                            If PATH is absolute, it is used as-is inside the container.
                            if PATH is relative, it is mounted relative to the project root
                            inside the container.
                            If NAME is not given, the volume name is generated.

  --mount-readonly HOST_DIR:CONTAINER_DIR
                            Mount a host file or directory to a container path as read-only.

  --mount-writable HOST_DIR:CONTAINER_DIR
                            Mount a host file or directory to a container path as writable.

  --publish [HOST_ADDRESS:]HOST_PORT:CONTAINER_PORT
                            Publish container port(s) to the host.
                            If no host address is given, ports bind to 127.0.0.1 by default.

  --tty                     Allocate a pseudo-TTY for the container.
  --no-cache                Disables cache when building the image.
  --skip-build              Skips building the image; assumes it already exists.
  --verbose                 Output verbose logs to stderr.
  --dry-run                 Does not execute the command; just prints it to stdout.
  --keep-alive SECONDS      Keep the container running after COMMAND finishes, so it can be reused.
                            Stop the container if no command is running for SECONDS.
                            Default: 0 (stop within ~5s after execution).


Examples:

  Start shell with preset configuration:
    $SCRIPT_NAME --tty --verbose zsh

  Check preset configuration:
    $SCRIPT_NAME --tty --verbose --dry-run zsh

  Install tools with mise:
    $SCRIPT_NAME --allow-write --allow-net mise-versions.jdx.dev,nodejs.org \\
      --verbose mise use node@lts

  Use volume:
    $SCRIPT_NAME --volume $SCRIPT_NAME--global--home-npm:/home/node/.npm \\
      --volume node_modules \\
      --verbose npm install

  Allow access to docker host:
    $SCRIPT_NAME --allow-net host.docker.internal:3000 \\
      --verbose busybox nc host.docker.internal 3000 < /dev/null

  Run with Dockerfile:
    $SCRIPT_NAME --dockerfile Dockerfile.minimum --tty --verbose bash


Preset Configuration:

  When --dockerfile is not specified, a preset Debian stable image is used with:
  - System packages: busybox, bash, zsh (with grml config), ripgrep, fd, dig, curl, git
  - mise package manager for additional runtime installations
  - Persistent storage for shell history, git config
  - Default editor: busybox vi


How to view DNS query log:

  DNS queries are logged by dnsmasq and can only be viewed through Docker logs.
  
  Find the container name:
  docker ps | grep $SCRIPT_NAME
  
  View DNS query logs in real-time
  docker logs -f <container-name>

Custom entrypoint:
  If /sandbox/bin/user-entrypoint.sh exists and is executable,
  it will be executed before the main command. Use this for custom
  initialization (e.g., setting environment variables, starting services).
HELP
}

log() {
  echo -e "[$(date "+%Y-%m-%d %H:%M:%S%z")][$(hostname)] $*"
}

# These variables are global because they are accessed in the on_exit trap after main() returns.
VERBOSE="no"
DRY_RUN="no"
IMAGE_NAME=

main() {
  # Options
  local dockerfile=""
  local allow_write="no"
  local allow_net="no"
  local allow_net_destinations=()
  local volume_dirs=()
  local readonly_mounts=()
  local writable_mounts=()
  local publish_ports=()
  local env_file=""
  local allocate_tty="no"
  local no_cache="no"
  local skip_build="no"
  local platform=""
  local keep_alive_seconds="0"
 
  while test "$#" -gt 0; do
    case "$1" in
      --help        ) help; return 0 ;;
      --dockerfile  ) dockerfile="$2"; shift 2 ;;
      --platform    ) platform="$2"; shift 2 ;;
      --allow-write ) allow_write="yes"; shift ;;
      --allow-net   )
        allow_net="yes"
        if test -n "$2" && ! grep -q "^--" <<< "$2"; then
          local new_allow_net_destinations
          IFS=',' read -r -a new_allow_net_destinations <<< "$2"
          allow_net_destinations+=("${new_allow_net_destinations[@]}")
          shift 2
        else
          shift 1
        fi
        ;;
      --volume )
        local new_volume_dirs
        IFS=',' read -r -a new_volume_dirs <<< "$2"
        volume_dirs+=("${new_volume_dirs[@]}")
        shift 2
        ;;
      --mount-readonly )
        local new_readonly_mounts
        IFS=',' read -r -a new_readonly_mounts <<< "$2"
        readonly_mounts+=("${new_readonly_mounts[@]}")
        shift 2
        ;;
      --mount-writable )
        local new_writable_mounts
        IFS=',' read -r -a new_writable_mounts <<< "$2"
        writable_mounts+=("${new_writable_mounts[@]}")
        shift 2
        ;;
      --publish )
        local new_publish_ports
        IFS=',' read -r -a new_publish_ports <<< "$2"
        publish_ports+=("${new_publish_ports[@]}")
        shift 2
        ;;
      --env-file   ) env_file="$2"; shift 2 ;;
      --tty        ) allocate_tty="yes"; shift ;;
      --no-cache   ) no_cache="yes"; shift ;;
      --skip-build ) skip_build="yes"; shift ;;
      --verbose    ) VERBOSE="yes"; shift ;;
      --dry-run    ) DRY_RUN="yes"; shift ;;
      --keep-alive ) keep_alive_seconds="$2"; shift 2 ;;
      --  ) shift; break ;;
      --* ) echo "Error: Unknown option: $1" >&2; return 1 ;;
      *   ) break ;;
    esac
  done

  # Save arguments to generate container id
  local sandbox_args=(
    "$dockerfile"
    "$allow_write"
    "$allow_net"
    "${allow_net_destinations[@]:-}"
    "${volume_dirs[@]:-}"
    "${readonly_mounts[@]:-}"
    "${writable_mounts[@]:-}"
    "${publish_ports[@]:-}"
    "$env_file"
    "$platform"
  )
  
  if test "$#" -eq 0; then
    help >&2
    return 1
  fi

  # save stdout, stderr
  exec 3>&1
  exec 4>&2

  if test "$VERBOSE" = "no"; then
    # discard stderr
    exec 2> /dev/null
  fi

  # stdout to stderr
  exec 1>&2

  run() {
    if test "$DRY_RUN" = "no"; then
      "$@"
    else
      echo "DRY_RUN: $*" >&3
    fi
  }

  log "Copy script to mountable path: $SCRIPT_PATH -> $MOUNTABLE_SCRIPT_PATH"
  run mkdir -p "$(dirname "$MOUNTABLE_SCRIPT_PATH")"
  run cp -f "$SCRIPT_PATH" "$MOUNTABLE_SCRIPT_PATH"
  run chmod +x "$MOUNTABLE_SCRIPT_PATH"

  local host_user_id
  local host_group_id
  host_user_id=$(id -u)
  host_group_id=$(id -g)
  
  local host_project_root
  local container_project_root
  local container_workdir
  host_project_root="$(git rev-parse --show-toplevel || pwd)"
  container_project_root="$host_project_root"
  container_workdir=$(pwd)

  local project_id
  if which shasum &> /dev/null; then
    project_id="$(shasum -a 256 <<< "$(pwd)" | head -c 8)"
  elif which sha256sum &> /dev/null; then
    project_id="$(sha256sum <<< "$(pwd)" | head -c 8)"
  else
    echo "Error: Neither shasum nor sha256sum found. Please install one of them." >&2
    return 1
  fi

  IMAGE_NAME="${SCRIPT_NAME}--$(basename "$(pwd)")-$project_id"

  local image_tag
  image_tag="latest"

  local container_id
  if which shasum &> /dev/null; then
    container_id="$(shasum -a 256 <<< "${sandbox_args[@]}" | head -c 8)"
  elif which sha256sum &> /dev/null; then
    container_id="$(sha256sum <<< "${sandbox_args[@]}" | head -c 8)"
  else
    echo "Error: Neither shasum nor sha256sum found. Please install one of them." >&2
    return 1
  fi

  local container_name="${IMAGE_NAME}--${container_id}"
  local network_name="${container_name}"

  if test -z "$dockerfile"; then
    log "Dockerfile not specified, using preset configuration."
    "$SCRIPT_PATH" print_preset_dockerfile

    local platform_with_default="default"
    if test -n "$platform"; then
      platform_with_default="$platform"
    fi

    volume_dirs+=(
      # global persistent volume
      "${SCRIPT_NAME}--global--${platform_with_default}--mise-data:/persistent/mise-data"
      # project local persistent volume
      "${IMAGE_NAME}--home:/persistent/home"
    )
  fi

  # docker options
  local docker_build_opts=(--tag "$IMAGE_NAME:$image_tag")
  local docker_run_opts=(
    --detach
    --rm
    --name "$container_name"
    --entrypoint ""
    --user 0:0
    --mount "type=bind,source=${MOUNTABLE_SCRIPT_PATH},target=${CONTAINER_SCRIPT_PATH},readonly"
  )
  local docker_exec_opts=(
    --interactive
    --user "$host_user_id:$host_group_id"
    --workdir "$container_workdir"
  )

  local docker_build_context=""
  if test -n "$dockerfile"; then
    docker_build_opts+=(--file "$dockerfile")
    docker_build_context=$(dirname "$dockerfile")
  fi

  if test "$no_cache" = "yes"; then
    docker_build_opts+=(--no-cache)
  fi

  if test -n "$platform"; then
    docker_build_opts+=(--platform "$platform")
    docker_run_opts+=(--platform "$platform")
  fi
  
  if test -n "$env_file"; then
    docker_run_opts+=(--env-file "$env_file")
  fi
  
  if test "$allow_write" = "yes"; then
    docker_run_opts+=(--mount "type=bind,source=${host_project_root},target=${container_project_root},consistency=delegated")
  else
    docker_run_opts+=(--mount "type=bind,source=${host_project_root},target=${container_project_root},readonly,consistency=delegated")
  fi
  
  if test "$allow_net" = "yes"; then
    docker_run_opts+=(
      --cap-add NET_ADMIN
      --cap-add NET_RAW
      --net "$network_name"
      --add-host host.docker.internal:host-gateway
    )
  else
    docker_run_opts+=(--net none)
  fi
  
  # volume options
  local volume_mount_paths=()
  if test "${#volume_dirs[@]}" -gt 0; then
    local volume_name
    local mount_path

    for dir in "${volume_dirs[@]}"; do
      if grep -qE ':' <<< "$dir"; then
        IFS=':' read -r volume_name mount_path <<< "$dir"
      else
        volume_name="${IMAGE_NAME}--$(echo "$dir" | sed 's,/,-,g; s,\.,-dot-,g')"
        mount_path="$dir"
        if ! grep -qE "^/" <<< "$mount_path"; then
          if ! test -e "$dir"; then
            mkdir -p "$dir"
          fi
          mount_path=$(readlink -f "$dir")
        fi
      fi

      docker_run_opts+=(--mount "type=volume,source=${volume_name},target=${mount_path},consistency=delegated")
      volume_mount_paths+=("$mount_path")
    done
  fi

  # mount options
  if test "${#readonly_mounts[@]}" -gt 0; then
    for mount in "${readonly_mounts[@]}"; do
      local host_path
      local container_path
      IFS=':' read -r host_path container_path <<< "$mount"
      local host_abs_path="$host_path"
      if ! grep -qE '^/' <<< "$host_abs_path"; then
        # shellcheck disable=SC2001
        host_abs_path=$(readlink -f "$(sed "s,~/,$HOME/," <<< "$host_abs_path")")
      fi
      docker_run_opts+=(--mount "type=bind,source=${host_abs_path},target=${container_path},readonly,consistency=delegated")
    done
  fi
  
  if test "${#writable_mounts[@]}" -gt 0; then
    for mount in "${writable_mounts[@]}"; do
      local host_path
      local container_path
      IFS=':' read -r host_path container_path <<< "$mount"
      local host_abs_path="$host_path"
      if ! grep -qE '^/' <<< "$host_abs_path"; then
        # shellcheck disable=SC2001
        host_abs_path=$(readlink -f "$(sed "s,~/,$HOME/," <<< "$host_abs_path")")
      fi
      docker_run_opts+=(--mount "type=bind,source=${host_abs_path},target=${container_path},consistency=delegated")
    done
  fi

  # publish options
  if test "${#publish_ports[@]}" -gt 0; then
    for port in "${publish_ports[@]}"; do
      if grep -qE '.+:.+:.+' <<< "$port"; then
        # host_address:host_port:container_port
        IFS=':' read -r host_address host_port container_port <<< "$port"
        docker_run_opts+=(--publish "${host_address}:${host_port}:${container_port}")
      elif grep -qE '.+:.+' <<< "$port"; then
        # host_port:container_port
        IFS=':' read -r host_port container_port <<< "$port"
        docker_run_opts+=(--publish "127.0.0.1:${host_port}:${container_port}")
      else
        echo "Error: Invalid port format: $port" >&2
        return 1
      fi
    done
  fi
  
  if test "$allocate_tty" = "yes"; then
    docker_exec_opts+=(--tty)
  fi

  local host_timezone
  if test -n "${TZ:-}"; then
    host_timezone="$TZ"
  elif readlink /etc/localtime | grep -q "/zoneinfo/"; then
    host_timezone="$(readlink /etc/localtime | awk -F '/' '{print $(NF-1)"/"$NF}')"
  elif test -f /etc/localtime; then
    host_timezone="$(cat /etc/timezone)"
  fi

  if test -n "$host_timezone"; then
    docker_run_opts+=(--env "TZ=${host_timezone}")
  fi

  # shellcheck disable=SC2001
  log "$(cat << EOF
Sandbox Configurations:
  docker_build_opts=
    $(echo "${docker_build_opts[*]}" | sed 's, ,\n    ,g')
  docker_build_context=$docker_build_context
  docker_run_opts=
    $(echo "${docker_run_opts[*]}" | sed 's, ,\n    ,g')
  docker_exec_opts=
    $(echo "${docker_exec_opts[*]}" | sed 's, ,\n    ,g')
  command=$@
EOF
  )"

  on_exit() {
    local exit_status="$?"
  
    if test "$VERBOSE" = "no"; then
      # discard stderr
      exec 2> /dev/null
    fi
    exec 1>&2
  
    run cleanup_networks "$IMAGE_NAME"
  
    exec 3>&-
    exec 4>&-
  
    return "$exit_status"
  }
 
  trap 'on_exit' EXIT

  if test "$skip_build" = "yes"; then
    log "Skip building docker image: $IMAGE_NAME"
  else
    log "Building docker image: $IMAGE_NAME"
    if test -n "$dockerfile"; then
      run docker build "${docker_build_opts[@]}" "$docker_build_context"
    else
      run docker build "${docker_build_opts[@]}" - \
        < <("$SCRIPT_PATH" print_preset_dockerfile)
    fi
  fi

  if test "$skip_build" = "yes" -a "$(docker inspect --type container "$container_name" --format "{{ .State.Running }}" 2> /dev/null)" = "true"; then
    log "Container is already running. Reusing $container_name"
  else
    log "Stopping any existing container: $container_name"
    run docker stop "$container_name" 2> /dev/null || true

    log "Remove any existing network: $network_name"
    run docker network rm "$network_name" 2> /dev/null || true

    if test "$allow_net" = "yes"; then
      log "Creating network: $network_name"
      run docker network create "$network_name" --driver bridge \
        -o com.docker.network.bridge.enable_ip_masquerade=true
    fi

    log "Inspect container user and group id."
    local container_default_user_group
    local container_default_user_id
    local container_default_group_id
    if test "$DRY_RUN" = "no"; then
      container_default_user_group=$(docker run --rm "${IMAGE_NAME}:${image_tag}" bash -c 'echo $(id -u):$(id -g)')
    else
      container_default_user_group="unknown:unknown"
    fi
    container_default_user_id=$(cut -d: -f1 <<< "$container_default_user_group")
    container_default_group_id=$(cut -d: -f2 <<< "$container_default_user_group")
    log "Container default user and group id: $container_default_user_group"

    log "Starting container."
    run docker run "${docker_run_opts[@]}" "${IMAGE_NAME}:${image_tag}" \
      "$CONTAINER_SCRIPT_PATH" start_container_dnsmasq

    if test "${#volume_mount_paths[@]}" -gt 0; then
      log "Setting up volume ownership to match host user ($host_user_id:$host_group_id)."
      run docker exec --user 0:0 "$container_name" "$CONTAINER_SCRIPT_PATH" setup_container_volume_owner \
        "$host_user_id" "$host_group_id" "${volume_mount_paths[@]}"
    fi

    if test "$allow_net" = "yes"; then
      log "Setting up firewall."
      run docker exec --user 0:0 "$container_name" \
        "$CONTAINER_SCRIPT_PATH" setup_container_firewall "${allow_net_destinations[@]}"
    fi

    log "Setting up container user to match host user ($host_user_id:$host_group_id) for file access."
    run docker exec --user 0:0 "$container_name" "$CONTAINER_SCRIPT_PATH" setup_container_user \
      "$container_default_user_id" "$container_default_group_id" \
      "$host_user_id" "$host_group_id"
  fi

  run docker exec --detach --user 0:0 "$container_name" \
    "$CONTAINER_SCRIPT_PATH" terminate_idle_container "$keep_alive_seconds"

  # restore stdout, stderr
  exec 1>&3
  exec 2>&4

  run docker exec "${docker_exec_opts[@]}" "$container_name" \
    "$CONTAINER_SCRIPT_PATH" container_exec "$container_project_root" "$@"
}

terminate_idle_container() {
  local keep_alive_seconds="$1"

  local last_activity_time
  last_activity_time=$(date +%s)

  # stop existing watchdog process
  for pid in $(busybox ps -o pid=,user=,comm=,args= | awk '$1 > 1 && $0 ~ /terminate_idle_container/ && $3 !~ /awk/ { print $1 }'); do
    if test "$pid" -ne "$$"; then
      kill "$pid" || true
    fi
  done

  while sleep 5; do
    if busybox ps -o pid=,user=,comm=,args= | awk '$1 > 1 && $2 !~ /root/ { print; found=1 } END { if (!found) exit 1 }'; then
      last_activity_time=$(date +%s)
    elif test "$(( $(date +%s) - last_activity_time ))" -ge "$keep_alive_seconds"; then
      kill 1
    fi
  done
}

cleanup_networks() {
  local image_name="$1"

  log "Cleaning up networks."
  local network_name
  for network_name in $(docker network ls --format '{{.Name}}' --filter "name=^$image_name"); do
    log "Checking if network $network_name is used."
    # Note: network_name = container_name
    if ! docker inspect --type container "$network_name" &> /dev/null; then
      log "Removing network."
      docker network rm "$network_name" || true
    fi
  done
}

container_exec() {
  local project_dir="$1"
  shift 1

  # Wait for the project directory owner to match the execution user (on colima).
  local max_attempts=20
  for _ in $(seq "$max_attempts"); do
    if stat -c "%U:%G" "$project_dir" | grep -qE "$(id -un):$(id -gn)"; then
      if test -x /sandbox/bin/user-entrypoint.sh; then
        exec /sandbox/bin/user-entrypoint.sh "$@"
      else
        exec "$@"
      fi
    fi
    sleep 0.5
  done

  echo "Error: Could not match project directory owner after multiple attempts." >&2
  exit 1
}

start_container_dnsmasq() {
  local servers=()
  IFS=$'\n' read -r -a servers < <(grep nameserver /etc/resolv.conf | awk '{print $2}')

  mkdir -p /sandbox/etc
  cat > /sandbox/etc/dnsmasq.conf << 'EOF'
log-queries
listen-address=127.0.0.1
EOF
  for server in "${servers[@]}"; do
    echo "server=$server" >> /sandbox/etc/dnsmasq.conf
  done
  echo "nameserver 127.0.0.1" > /etc/resolv.conf

  exec dnsmasq -k -C /sandbox/etc/dnsmasq.conf --log-facility /dev/stdout
}

setup_container_volume_owner() {
  local host_user_id=$1
  local host_group_id=$2
  shift 2

  for mount_path in "$@"; do
    chown "$host_user_id:$host_group_id" "$mount_path"
  done
}

setup_container_firewall() {
  local destinations=("$@")
  local addresses=()

  local address_pattern='^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}(/[0-9]{1,2})?$'
  
  for destination in "${destinations[@]}"; do
    local host="$destination"
    local port=''
    if grep -q ':' <<< "$destination"; then
      host=$(cut -d: -f1 <<< "$destination")
      port=$(cut -d: -f2 <<< "$destination")
    fi

    # address
    if grep -qE "$address_pattern" <<< "$host"; then
      echo "Allow outgoing connections to $destination"
      if test -n "$port"; then
        addresses+=("$host:$port")
      else
        addresses+=("$host")
      fi
      continue
    fi

    # domain
    local domain_address_count=0
    for address in $(getent hosts "$host" | awk '{print $1}'; dig +short A "$host"); do
      if ! grep -qE "$address_pattern" <<< "$address"; then
        log "Warning: Ignoring invalid address $address of $host"
        continue
      fi

      if test -n "$port"; then
        log "Allow outgoing connections to $destination $address:$port"
        addresses+=("$address:$port")
      else
        log "Allow outgoing connections to $destination $address"
        addresses+=("$address")
      fi

      domain_address_count=$((domain_address_count + 1))
    done

    if test "$domain_address_count" = 0; then
      echo "Error: Failed to resolve any address for $host" >&2
      return 1
    fi
  done
  
  local docker_host_address
  docker_host_address=$(busybox ip route | grep default | cut -d" " -f3)
  if test -z "$docker_host_address"; then
    echo "Error: Failed to determine docker host address" >&2
    return 1
  fi

  # reset
  iptables -F
  iptables -X
  # iptables -t nat -F
  # iptables -t nat -X
  iptables -t mangle -F
  iptables -t mangle -X

  ip6tables -F
  ip6tables -X
  ip6tables -t mangle -F
  ip6tables -t mangle -X

  ipset create allow_list hash:net,port -exist
  ipset flush allow_list
  
  # allow dns
  iptables -A OUTPUT -p udp --dport 53 -j ACCEPT
  iptables -A INPUT -p udp --sport 53 -j ACCEPT
  
  # allow localhost
  iptables -A INPUT -i lo -j ACCEPT
  iptables -A OUTPUT -o lo -j ACCEPT
  
  # set default policies
  iptables -P INPUT DROP
  iptables -P FORWARD DROP
  iptables -P OUTPUT DROP

  ip6tables -P INPUT DROP
  ip6tables -P FORWARD DROP
  ip6tables -P OUTPUT DROP
  
  # allow access from docker host
  iptables -A INPUT -s "$docker_host_address" -j ACCEPT
  # Do not allow access to docker host by default
  # iptables -A OUTPUT -d "$docker_host_address" -j ACCEPT
  
  # allow established connections
  iptables -A INPUT -m state --state ESTABLISHED,RELATED -j ACCEPT
  iptables -A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT

  # allow localhost on IPv6
  ip6tables -A INPUT -i lo -j ACCEPT
  ip6tables -A OUTPUT -o lo -j ACCEPT

  # allow established connections on IPv6
  ip6tables -A INPUT -m state --state ESTABLISHED,RELATED -j ACCEPT
  ip6tables -A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT
  
  # allow given addresses
  for address in "${addresses[@]}"; do
    local address_part
    local port
    if grep -q ':' <<< "$address"; then
      IFS=':' read -r address_part port <<< "$address"
    else
      address_part="$address"
      port=443
    fi

    if test "$address_part" = "0.0.0.0/0"; then
      iptables -A OUTPUT -d "$address_part" -p tcp --dport "$port" -j ACCEPT
    else
      ipset add allow_list "${address_part},tcp:${port}" -exist
    fi
  done
  iptables -A OUTPUT -p tcp -m set --match-set allow_list dst,dst -j ACCEPT

  # Reject unauthorized TCP connections immediately with RST
  # instead of silently dropping and causing timeout
  iptables -A OUTPUT -p tcp -j REJECT --reject-with tcp-reset
  ip6tables -A OUTPUT -p tcp -j REJECT --reject-with tcp-reset
}

setup_container_user() {
  local container_default_user_id="$1"
  local container_default_group_id="$2"
  local host_user_id="$3"
  local host_group_id="$4"

  local user
  local group
  user=$(getent passwd "$container_default_user_id" | cut -d: -f1 2> /dev/null || echo "")
  group=$(getent group "$container_default_group_id" | cut -d: -f1 2> /dev/null || echo "")

  if test "$container_default_user_id" -eq 0; then
    log "Container's default user is root, creating sandbox user/group"
    
    if ! getent group sandbox &> /dev/null; then
      log "Creating group sandbox"
      groupadd sandbox
    fi
    
    if ! getent passwd sandbox &> /dev/null; then
      log "Creating user sandbox"
      useradd -g sandbox -m sandbox
    fi
    
    user="sandbox"
    group="sandbox"
  fi

  log "Updating group '$group' ID to $host_group_id"
  local conflicting_group
  conflicting_group=$(getent group "$host_group_id" | cut -d: -f1 2> /dev/null || echo "")
  if test -n "$conflicting_group" && test "$conflicting_group" != "$group"; then
    log "Resolving group ID conflict: moving group '$conflicting_group'"
    local temp_gid
    for gid in $(seq 1000 65533); do
      if ! getent group "$gid" &> /dev/null; then
        temp_gid=$gid
        break
      fi
    done
    
    if test -z "$temp_gid"; then
      echo "Error: No available GID in user range" >&2
      return 1
    fi
    
    groupmod -g "$temp_gid" "$conflicting_group"
  fi

  groupmod -g "$host_group_id" "$group"

  log "Updating user '$user' ID to $host_user_id:$host_group_id"
  local conflicting_user
  conflicting_user=$(getent passwd "$host_user_id" | cut -d: -f1 2> /dev/null || echo "")
  if test -n "$conflicting_user" && test "$conflicting_user" != "$user"; then
    log "Resolving user ID conflict: moving user '$conflicting_user'"
    local temp_uid
    for uid in $(seq 1000 65533); do
      if ! getent passwd "$uid" &> /dev/null; then
        temp_uid=$uid
        break
      fi
    done
    
    if test -z "$temp_uid"; then
      echo "Error: No available UID in user range" >&2
      return 1
    fi
    
    usermod -u "$temp_uid" "$conflicting_user"
  fi

  usermod -u "$host_user_id" -g "$host_group_id" "$user"

  # Adjust home directory permissions to match new UID/GID
  local home_dir
  home_dir=$(getent passwd "$user" | cut -d: -f6)
  if test -n "$home_dir" && test -d "$home_dir"; then
    log "Updating home directory ownership: $home_dir"
    chown "$host_user_id:$host_group_id" "$home_dir" || true
  fi

  log "User setup completed: $user ($host_user_id:$host_group_id)"
}

print_preset_dockerfile() {
  echo "FROM $PRESET_BASE_IMAGE"
  cat << 'EOF'

RUN apt update \
    && apt install -y \
      busybox bash zsh locales gpg \
      fd-find ripgrep jq \
      iptables ipset dnsmasq dnsutils curl \
      build-essential git tmux \
    && bash -c 'ln -s $(which fdfind) /usr/local/bin/fd' \
    && echo 'en_US.UTF-8 UTF-8' >> /etc/locale.gen \
    && echo 'ja_JP.UTF-8 UTF-8' >> /etc/locale.gen \
    && locale-gen

# mise: https://mise.jdx.dev/
RUN install -dm 755 /etc/apt/keyrings \
    && curl -fsSL https://mise.jdx.dev/gpg-key.pub | gpg --dearmor | tee /etc/apt/keyrings/mise-archive-keyring.gpg 1> /dev/null \
    && echo "deb [signed-by=/etc/apt/keyrings/mise-archive-keyring.gpg arch=$(dpkg --print-architecture)] https://mise.jdx.dev/deb stable main" | tee /etc/apt/sources.list.d/mise.list \
    && apt update \
    && apt install -y mise

# Create user
RUN groupadd sandbox \
    && useradd -g sandbox -m sandbox \
    && mkdir -p /sandbox \
    && chmod 777 /sandbox

# Create user entrypoint script for persistence configuration
RUN mkdir -p /sandbox/bin
COPY <<'USER_ENTRYPOINT' /sandbox/bin/user-entrypoint.sh
#!/bin/bash
# Configure persistence for sandbox data (history, configs, etc.) by symlinking to /persistent/home/.

test -L ~/.bash_history || (touch /persistent/home/.bash_history; ln -s /persistent/home/.bash_history ~/ 2> /dev/null)
test -L ~/.zsh_history || (touch /persistent/home/.zsh_history; ln -s /persistent/home/.zsh_history ~/ 2> /dev/null)
test -L ~/.config || (mkdir -p /persistent/home/.config; ln -s /persistent/home/.config ~/ 2> /dev/null)
test -L ~/.gitconfig || (touch /persistent/home/.gitconfig; ln -s /persistent/home/.gitconfig ~/ 2> /dev/null)
test -L ~/.local || (mkdir -p /persistent/home/.local/share; ln -s /persistent/home/.local ~/ 2> /dev/null)
test -L ~/.gemini || (mkdir -p /persistent/home/.gemini; ln -s /persistent/home/.gemini ~/ 2> /dev/null)
test -L ~/.codex || (mkdir -p /persistent/home/.codex; ln -s /persistent/home/.codex ~/ 2> /dev/null)

eval "$(mise activate bash)"
exec "$@"
USER_ENTRYPOINT
RUN chmod +x /sandbox/bin/user-entrypoint.sh

USER sandbox

# Configure shell
# - grml zsh config: https://grml.org/zsh/
# - zsh-autosuggestions: https://github.com/zsh-users/zsh-autosuggestions
# - zsh-syntax-highlighting: https://github.com/zsh-users/zsh-syntax-highlighting
RUN curl -fsSL https://raw.githubusercontent.com/grml/grml-etc-core/refs/tags/v0.19.23/etc/zsh/zshrc -o /home/sandbox/.zshrc \
    && echo 'fc0a6642d61193fd95293fb39a561e456e240b81e9ddb8c08fc5c2ac8c31d2f4 /home/sandbox/.zshrc' > /home/sandbox/.zshrc.sha256sum \
    && sha256sum -c /home/sandbox/.zshrc.sha256sum \
    && git clone --depth 1 https://github.com/zsh-users/zsh-autosuggestions /home/sandbox/.zsh/zsh-autosuggestions \
    && echo 'source ~/.zsh/zsh-autosuggestions/zsh-autosuggestions.zsh' >> /home/sandbox/.zshrc.local \
    && git clone --depth 1 https://github.com/zsh-users/zsh-syntax-highlighting /home/sandbox/.zsh/zsh-syntax-highlighting \
    && echo 'source ~/.zsh/zsh-syntax-highlighting/zsh-syntax-highlighting.zsh' >> /home/sandbox/.zshrc.local \
    && echo 'unsetopt HIST_SAVE_BY_COPY' >> /home/sandbox/.zshrc.local

ENV LANG=en_US.UTF-8
ENV EDITOR="busybox vi"
ENV MISE_DATA_DIR=/persistent/mise-data
EOF
}

if test "$#" -gt 0; then
  case "$1" in
    start_container_dnsmasq \
      | setup_container_volume_owner \
      | setup_container_firewall \
      | setup_container_user \
      | terminate_idle_container \
      | container_exec \
      | print_preset_dockerfile )
      "$@"
      ;;
    * )
      main "$@"
  esac
else
  main
fi
