#!/bin/sh
set -eu

usage() {
  echo "usage: $0 <service> <repo> <main> <image> <cluster> <arch> <work-dir> [mode] [push-profile] [push-region]" >&2
  exit 2
}

[ "$#" -ge 7 ] || usage

service=$1
repo=$2
main_path=$3
image=$4
cluster=$5
arch=$6
work_dir=$7
mode=${8:-generic}
push_profile=${9:-}
push_region=${10:-}

case "$mode" in
  generic | generic-cgo | serverless-global | serverless-regional | serverless-service) ;;
  *)
    echo "unsupported Starter build mode: $mode" >&2
    exit 2
    ;;
esac

repo=$(cd "$repo" && pwd)
script_dir=$(cd "$(dirname "$0")" && pwd)
mkdir -p "$work_dir"
work_dir=$(cd "$work_dir" && pwd)
live_binary_output=${CP9_LIVE_BINARY_OUTPUT:-}
enable_live_update=${CP9_ENABLE_LIVE_UPDATE:-}
ldflags_arg=
cgo_enabled=0
if [ "$mode" = "generic-cgo" ]; then
  if [ "$(uname -s)" != "Linux" ]; then
    echo "$service requires the Linux remote builder for its CGO dependency" >&2
    exit 1
  fi
  cgo_enabled=1
fi
if [ -n "$push_profile" ]; then
  ldflags_arg="-s -w"
fi

if [ "$mode" = "serverless-regional" ]; then
  source_root=$(dirname "$repo")
  service_dir=$(basename "$repo")
  patched_repo=$work_dir/source
  rm -rf "$patched_repo"
  mkdir -p "$patched_repo"
  (cd "$source_root" && tar --exclude .git -cf - .) | (cd "$patched_repo" && tar -xf -)
  repo=$patched_repo/$service_dir
fi

if [ "$mode" = "serverless-global" ]; then
  source_root=$(dirname "$repo")
  service_dir=$(basename "$repo")
  patched_repo=$work_dir/source
  rm -rf "$patched_repo"
  mkdir -p "$patched_repo"
  (cd "$source_root" && tar --exclude .git -cf - .) | (cd "$patched_repo" && tar -xf -)
  repo=$patched_repo/$service_dir
fi

goflags=${GOFLAGS:-}
case " $goflags " in
  *" -buildvcs="*) ;;
  *)
    if [ -n "$goflags" ]; then
      goflags="$goflags -buildvcs=false"
    else
      goflags="-buildvcs=false"
    fi
    ;;
esac

main_dir=$(dirname "$main_path")
build_root=$repo
if [ "$main_dir" != "." ]; then
  current=$repo/$main_dir
  while [ "$current" != "$repo" ] && [ "$current" != "/" ]; do
    if [ -f "$current/go.mod" ]; then
      build_root=$current
      break
    fi
    current=$(dirname "$current")
  done
fi
if [ ! -f "$build_root/go.mod" ]; then
  build_root=$repo
fi

rel_pkg=.
case "$main_path" in
  */*) rel_pkg=./$(dirname "$main_path") ;;
esac
case "$build_root" in
  "$repo")
    ;;
  "$repo"/*)
    rel_root=${build_root#"$repo"/}
    rel_pkg=./$(dirname "$main_path")
    rel_pkg=${rel_pkg#./$rel_root/}
    rel_pkg=./$rel_pkg
    ;;
esac

rootfs=$work_dir/rootfs
dockerfile=$work_dir/Dockerfile
rm -rf "$rootfs"
mkdir -p "$rootfs/usr/local/bin"

entrypoint=/usr/local/bin/service
binary=$work_dir/service
printf '%s\n' "[cp9 persona build] $service: go build $rel_pkg"
if [ -f "$build_root/go.mod" ]; then
  if [ -n "$ldflags_arg" ]; then
    GOPRIVATE=github.com/tidbcloud/* GONOSUMDB=github.com/tidbcloud/* \
      GOFLAGS="$goflags" \
      CGO_ENABLED="$cgo_enabled" GOOS=linux GOARCH="$arch" \
      go -C "$build_root" build -ldflags "$ldflags_arg" -o "$binary" "$rel_pkg"
  else
    GOPRIVATE=github.com/tidbcloud/* GONOSUMDB=github.com/tidbcloud/* \
      GOFLAGS="$goflags" \
      CGO_ENABLED="$cgo_enabled" GOOS=linux GOARCH="$arch" \
      go -C "$build_root" build -o "$binary" "$rel_pkg"
  fi
else
  if [ -n "$ldflags_arg" ]; then
    GO111MODULE=off CGO_ENABLED="$cgo_enabled" GOOS=linux GOARCH="$arch" \
      go -C "$build_root" build -ldflags "$ldflags_arg" -o "$binary" "$rel_pkg"
  else
    GO111MODULE=off CGO_ENABLED="$cgo_enabled" GOOS=linux GOARCH="$arch" \
      go -C "$build_root" build -o "$binary" "$rel_pkg"
  fi
fi

cp "$binary" "$rootfs/usr/local/bin/service"
chmod 755 "$rootfs/usr/local/bin/service"
mkdir -p "$rootfs/workspace"
mkdir -p "$rootfs/command/$service"
if [ "$enable_live_update" = "1" ]; then
  cp "$binary" "$rootfs/workspace/service"
  chmod 755 "$rootfs/workspace/service"
  cp "$binary" "$rootfs/command/$service/$service"
  chmod 755 "$rootfs/command/$service/$service"
else
  ln -s /usr/local/bin/service "$rootfs/workspace/service"
  ln -s /usr/local/bin/service "$rootfs/command/$service/$service"
fi
if [ "$service" = "serverless-service" ] && [ "$enable_live_update" != "1" ]; then
  ln -s /usr/local/bin/service "$rootfs/serverless-service"
fi
if [ "$service" = "cluster-proxy" ]; then
  ln -s /usr/local/bin/service "$rootfs/proxy"
fi
if [ "$service" = "gateway-controller" ]; then
  ln -s /usr/local/bin/service "$rootfs/gateway-controller"
fi
if [ "$service" = "serverless-meta-service" ]; then
  ln -s /usr/local/bin/service "$rootfs/serverless-metering"
fi
if [ "$service" = "vpa-agent" ]; then
  ln -s /usr/local/bin/service "$rootfs/vpa-agent"
fi

if [ -n "$live_binary_output" ]; then
  live_output_dir=$(dirname "$live_binary_output")
  mkdir -p "$live_output_dir"
  tmp_live=$live_binary_output.tmp.$$
  rm -f "$tmp_live"
  cp "$binary" "$tmp_live"
  chmod 755 "$tmp_live"
  mv "$tmp_live" "$live_binary_output"
  exit 0
fi

write_live_update_entrypoint_wrapper() {
  wrapper=$1
  mkdir -p "$(dirname "$rootfs$wrapper")"
  cat >"$rootfs$wrapper" <<'EOF'
#!/bin/sh
set -u

marker=/tmp/cp9-live-update.restart
child=0
stopping=0

marker_stamp() {
  if [ -e "$marker" ]; then
    stat -c %Y "$marker" 2>/dev/null || echo 0
    return
  fi
  echo 0
}

stop_child() {
  stopping=1
  if [ "$child" != 0 ]; then
    kill "$child" 2>/dev/null || true
  fi
}

trap stop_child INT TERM

last_stamp=$(marker_stamp)
while :; do
  /workspace/service "$@" &
  child=$!

  while kill -0 "$child" 2>/dev/null; do
    if [ "$stopping" = 1 ]; then
      wait "$child"
      exit $?
    fi
    current_stamp=$(marker_stamp)
    if [ "$current_stamp" != "$last_stamp" ]; then
      last_stamp=$current_stamp
      kill "$child" 2>/dev/null || true
      wait "$child" 2>/dev/null || true
      child=0
      break
    fi
    sleep 1
  done

  if [ "$child" != 0 ]; then
    wait "$child"
    exit $?
  fi
done
EOF
  chmod 755 "$rootfs$wrapper"
}

if [ "$enable_live_update" = "1" ]; then
  write_live_update_entrypoint_wrapper /usr/local/bin/service
  write_live_update_entrypoint_wrapper "/command/$service/$service"
  if [ "$service" = "serverless-service" ]; then
    write_live_update_entrypoint_wrapper /serverless-service
  fi
fi

cat >"$dockerfile" <<EOF
FROM public.ecr.aws/docker/library/alpine:3.19
COPY rootfs/ /
RUN apk add --no-cache ca-certificates && update-ca-certificates && mkdir -p /workspace && chown -R 65532:65532 /workspace
WORKDIR /workspace
USER 65532:65532
ENTRYPOINT ["$entrypoint"]
EOF

write_image_stamp() {
  stamp_tmp=$work_dir/image.stamp.tmp.$$
  image_id=$(docker image inspect "$image" --format '{{.Id}}' 2>/dev/null || true)
  if [ -z "$image_id" ]; then
    image_id=unknown-image
  fi
  build_time=$(date -u +%Y%m%dT%H%M%SZ)
  printf '%s\n' "$image_id $build_time" >"$stamp_tmp"
  mv "$stamp_tmp" "$work_dir/image.stamp"
  printf '%s\n' "[cp9 persona build] $service: wrote image rollout stamp $work_dir/image.stamp"
}

if [ -n "$push_profile" ]; then
  printf '%s\n' "[cp9 persona build] $service: docker build for remote ECR $image"
  docker_buildkit_config=
  buildx_available=0
  if docker buildx version >/dev/null 2>&1; then
    buildx_available=1
  else
    login_home=$(eval "printf '%s' ~$(id -un)")
    if [ -x "$login_home/.docker/cli-plugins/docker-buildx" ]; then
      if DOCKER_CONFIG="$login_home/.docker" docker buildx version >/dev/null 2>&1; then
        docker_buildkit_config=$login_home/.docker
        buildx_available=1
      fi
    fi
  fi
  if [ "$buildx_available" -eq 1 ]; then
    if [ -n "$docker_buildkit_config" ]; then
      DOCKER_CONFIG="$docker_buildkit_config" DOCKER_BUILDKIT=1 docker buildx build --pull --no-cache --provenance=false --load --platform="linux/$arch" -t "$image" -f "$dockerfile" "$work_dir"
    else
      DOCKER_BUILDKIT=1 docker buildx build --pull --no-cache --provenance=false --load --platform="linux/$arch" -t "$image" -f "$dockerfile" "$work_dir"
    fi
  else
    DOCKER_BUILDKIT=1 docker build --pull --no-cache --platform="linux/$arch" -t "$image" -f "$dockerfile" "$work_dir"
  fi
else
  printf '%s\n' "[cp9 persona build] $service: docker build $image"
  docker build --platform="linux/$arch" -t "$image" -f "$dockerfile" "$work_dir"
fi

if [ -n "$push_profile" ]; then
  [ -n "$push_region" ] || {
    echo "push region is required when push profile is set" >&2
    exit 2
  }
  registry=${image%%/*}
  push_attempts=${CP9_ECR_PUSH_ATTEMPTS:-10}
  push_sleep_seconds=${CP9_ECR_PUSH_RETRY_SLEEP_SECONDS:-5}
  case "$push_attempts" in
    "" | *[!0-9]*)
      echo "CP9_ECR_PUSH_ATTEMPTS must be a positive integer" >&2
      exit 2
      ;;
  esac
  case "$push_sleep_seconds" in
    "" | *[!0-9]*)
      echo "CP9_ECR_PUSH_RETRY_SLEEP_SECONDS must be a non-negative integer" >&2
      exit 2
      ;;
  esac
  if [ "$push_attempts" -lt 1 ]; then
    echo "CP9_ECR_PUSH_ATTEMPTS must be a positive integer" >&2
    exit 2
  fi

  run_crane() {
    if [ -n "${CP9_CRANE_BIN:-}" ]; then
      "$CP9_CRANE_BIN" "$@"
      return
    fi
    if command -v crane >/dev/null 2>&1; then
      crane "$@"
      return
    fi
    go run github.com/google/go-containerregistry/cmd/crane@v0.20.6 "$@"
  }
  crane_tar=$work_dir/crane-image.tar
  crane_config=$work_dir/crane-docker-config
  mkdir -p "$crane_config"
  docker save "$image" -o "$crane_tar"
  python3 - "$crane_tar" <<'PY'
import gzip
import hashlib
import json
import sys
import tarfile

archive_path = sys.argv[1]

with tarfile.open(archive_path) as archive:
    manifest = json.load(archive.extractfile("manifest.json"))
    if len(manifest) != 1:
        raise SystemExit(f"expected one image in docker archive, found {len(manifest)}")

    image = manifest[0]
    config = json.load(archive.extractfile(image["Config"]))
    expected = config.get("rootfs", {}).get("diff_ids", [])
    layers = image.get("Layers", [])
    if len(expected) != len(layers):
        raise SystemExit(f"archive layer count mismatch: config={len(expected)} manifest={len(layers)}")

    for index, (diff_id, layer_name) in enumerate(zip(expected, layers), start=1):
        layer = archive.extractfile(layer_name).read()
        if layer.startswith(b"\x1f\x8b"):
            layer = gzip.decompress(layer)
        actual = "sha256:" + hashlib.sha256(layer).hexdigest()
        if actual != diff_id:
            raise SystemExit(
                f"archive layer {index} diff id mismatch: config={diff_id} actual={actual}"
            )
PY
  pushed=0
  push_attempt=1
  while [ "$push_attempt" -le "$push_attempts" ]; do
    printf '%s\n' "[cp9 persona build] $service: docker push $image (attempt $push_attempt/$push_attempts)"
    if aws --profile "$push_profile" --region "$push_region" ecr get-login-password |
      DOCKER_CONFIG="$crane_config" docker login --username AWS --password-stdin "$registry" &&
      DOCKER_CONFIG="$crane_config" docker push "$image"; then
      pushed=1
      break
    fi
    if [ "$push_attempt" -eq "$push_attempts" ]; then
      break
    fi
    printf '%s\n' "[cp9 persona build] $service: docker push failed; retrying in ${push_sleep_seconds}s"
    sleep "$push_sleep_seconds"
    push_attempt=$((push_attempt + 1))
    if [ "$push_sleep_seconds" -gt 0 ]; then
      push_sleep_seconds=$((push_sleep_seconds * 2))
      if [ "$push_sleep_seconds" -gt 30 ]; then
        push_sleep_seconds=30
      fi
    fi
  done
  if [ "$pushed" -ne 1 ]; then
    echo "docker push failed after $push_attempts attempts: $image" >&2
    exit 1
  fi
fi

if [ -n "$cluster" ]; then
  printf '%s\n' "[cp9 persona build] $service: k3d import $image into $cluster"
  k3d image import "$image" -c "$cluster"
fi

write_image_stamp