#!/bin/bash

function klabels {
	local header=
	local padding=
	local trim_length=
	local line_text=
	local first_line=
	local max_key_length=0
	local max_label_length=0
	local key=
	local value=
	local filter_value="$1"
	local filter_key="$2"

	while IFS= read -r line; do
		if [ -z "$header" ]; then
			header="$line"
			padding="${header//LABELS/}"
			padding="${padding//?/ }"
			echo "$header"
			trim_length=${#padding}
			continue
		fi

		line_text="${line:0:trim_length}"
		first_line=
		IFS=',' read -ra labels <<< "${line:trim_length}"

		# Calculate the maximum key length to align key=value pairs
		for label in "${labels[@]}"; do
			key="${label%%=*}"
			max_key_length=$(( ${#key} > max_key_length ? ${#key} : max_key_length ))
		done

		# Loop through labels and align them
		for label in "${labels[@]}"; do
			key="${label%%=*}"
			value="${label#*=}"

			if [ -n "$filter_key" ] && ! [[ "$key" == *$filter_key* ]] ; then
				continue
			fi

			if [ -n "$filter_value" ] && ! [[ "$value" == *$filter_value* ]] ; then
				continue
			fi


			# Format the label to align the '=' sign
			printf -v formatted_label "%-${max_key_length}s = %s" "$key" "$value"

			if [ -z "$first_line" ]; then
				first_line=1

				# Print a line separator based on the longest key length
				separator_length=$((max_key_length + trim_length))
				printf "%${separator_length}s\n" | tr ' ' "-"

				echo "$line_text$formatted_label"
			else
				echo "$padding$formatted_label"
			fi
		done
	done
}

function __kcreate_help {
	cat <<'EOH'
usage:
  kcreate [-n NAMESPACE] [options] TARGET_POD NEW_POD [-exec [--] CMD...]

examples:
  kcreate -n www app-pod debug-pod
  kcreate -n www app-pod debug-pod -exec -- bash -l
  kcreate -n www -c app --rm app-pod debug-pod -exec -- bash -l

options:
  -n, --namespace NS       Namespace. Defaults to current kubectl namespace, then default.
  -c, --container NAME     Container to wait for / exec into.
  -e, -exec, --exec        After creating the Pod, wait for container(s) to run, then kubectl exec -it.
  --rm                    Delete the new Pod after the exec command exits.
  --sleep, --sleep-infinity
                          Start copied containers with: sleep infinity.
  --force                 Delete NEW_POD first if it already exists.
  --dry-run               Print the generated Pod manifest instead of creating it.
  --copy-labels           Copy labels from the target Pod. Dangerous: Services may route traffic to it.
  --copy-annotations      Copy annotations from the target Pod. Usually not needed.
  --same-node             Keep spec.nodeName, forcing the new Pod onto the same node.
  --no-probes             Remove liveness/readiness/startup probes from copied containers.
  --timeout SECONDS       Exec wait timeout. Default: 90.
EOH
}

function __kcreate_current_namespace {
	local ns

	ns=$(kubectl config view --minify -o 'jsonpath={..namespace}' 2>/dev/null)
	if [ -z "$ns" ]; then
		ns=default
	fi

	printf '%s\n' "$ns"
}

function __kcreate_wait_container_running {
	local ns="$1"
	local pod="$2"
	local container="$3"
	local timeout="$4"
	local end
	local ok

	end=$((SECONDS + timeout))

	while [ "$SECONDS" -le "$end" ]; do
		if [ -n "$container" ]; then
			ok=$(
				kubectl -n "$ns" get pod "$pod" -o json 2>/dev/null | jq -r --arg c "$container" '
					[.status.containerStatuses[]? | select(.name == $c and (.state.running != null))] | length > 0
				'
			) || ok=false
		else
			ok=$(
				kubectl -n "$ns" get pod "$pod" -o json 2>/dev/null | jq -r '
					[.status.containerStatuses[]?] as $s
					| ($s | length > 0) and ([ $s[] | select(.state.running != null) ] | length == ($s | length))
				'
			) || ok=false
		fi

		if [ "$ok" = true ]; then
			return 0
		fi

		sleep 1
	done

	echo "kcreate: timed out waiting for pod/$pod container(s) to be running" >&2
	kubectl -n "$ns" get pod "$pod" >&2
	return 1
}

function kcreate {
	local ns=
	local target_pod=
	local new_pod=
	local container=
	local sleep_forever=false
	local exec_after=false
	local rm_after=false
	local force=false
	local dry_run=false
	local copy_labels=false
	local copy_annotations=false
	local same_node=false
	local no_probes=false
	local timeout=90
	local manifest
	local rc
	local exec_cmd=()
	local exec_args=()

	while [ "$#" -gt 0 ]; do
		case "$1" in
			-h|--help)
				__kcreate_help
				return 0
				;;
			-n|--namespace)
				if [ -z "$2" ]; then echo "kcreate: $1 needs a namespace" >&2; return 2; fi
				ns="$2"
				shift 2
				;;
			-c|--container)
				if [ -z "$2" ]; then echo "kcreate: $1 needs a container name" >&2; return 2; fi
				container="$2"
				shift 2
				;;
			--sleep|--sleep-infinity)
				sleep_forever=true
				shift
				;;
			-e|-exec|--exec)
				exec_after=true
				shift
				if [ "${1:-}" = "--" ]; then shift; fi
				exec_cmd=("$@")
				break
				;;
			--rm)
				rm_after=true
				shift
				;;
			--force)
				force=true
				shift
				;;
			--dry-run)
				dry_run=true
				shift
				;;
			--copy-labels)
				copy_labels=true
				shift
				;;
			--copy-annotations)
				copy_annotations=true
				shift
				;;
			--same-node)
				same_node=true
				shift
				;;
			--no-probes)
				no_probes=true
				shift
				;;
			--timeout)
				if [ -z "$2" ]; then echo "kcreate: $1 needs seconds" >&2; return 2; fi
				timeout="$2"
				shift 2
				;;
			-*)
				echo "kcreate: unknown option: $1" >&2
				return 2
				;;
			*)
				if [ -z "$target_pod" ]; then
					target_pod="$1"
				elif [ -z "$new_pod" ]; then
					new_pod="$1"
				else
					echo "kcreate: unexpected argument: $1" >&2
					return 2
				fi
				shift
				;;
		esac
	done

	if [ -z "$ns" ]; then
		ns=$(__kcreate_current_namespace)
	fi

	if [ -z "$target_pod" ] || [ -z "$new_pod" ]; then
		__kcreate_help >&2
		return 2
	fi

	if ! command -v kubectl >/dev/null 2>&1; then
		echo "kcreate: kubectl is required" >&2
		return 127
	fi

	if ! command -v jq >/dev/null 2>&1; then
		echo "kcreate: jq is required" >&2
		return 127
	fi

	if ! [[ "$timeout" =~ ^[0-9]+$ ]]; then
		echo "kcreate: --timeout must be plain seconds, e.g. 90" >&2
		return 2
	fi

	manifest=$(
		kubectl -n "$ns" get pod "$target_pod" -o json | jq \
			--arg name "$new_pod" \
			--arg src "$target_pod" \
			--arg srcns "$ns" \
			--argjson sleep_forever "$sleep_forever" \
			--argjson copy_labels "$copy_labels" \
			--argjson copy_annotations "$copy_annotations" \
			--argjson same_node "$same_node" \
			--argjson no_probes "$no_probes" '
			{
				apiVersion: "v1",
				kind: "Pod",
				metadata: {
					name: $name,
					labels: (if $copy_labels then (.metadata.labels // {}) else {"app.kubernetes.io/managed-by": "kcreate"} end),
					annotations: (if $copy_annotations then (.metadata.annotations // {}) else {} end)
				},
				spec: .spec
			}
			| .metadata.annotations["kcreate/source-pod"] = $src
			| .metadata.annotations["kcreate/source-namespace"] = $srcns
			| del(.metadata.labels["pod-template-hash"]?)
			| del(.spec.priority?)
			| del(.spec.serviceAccount?)
			| if $same_node then . else del(.spec.nodeName?) end
			| if $no_probes then
				(.spec.containers |= map(del(.livenessProbe?, .readinessProbe?, .startupProbe?)))
			  else
				.
			  end
			| if $sleep_forever then
				(.spec.containers |= map(
					.command = ["sleep"]
					| .args = ["infinity"]
					| del(.livenessProbe?, .readinessProbe?, .startupProbe?)
				))
			  else
				.
			end
		'
	) || return $?

	if [ "$dry_run" = true ]; then
		printf '%s\n' "$manifest"
		return 0
	fi

	if [ "$copy_labels" = true ]; then
		echo "kcreate: warning: --copy-labels may let Services route traffic to pod/$new_pod" >&2
	fi

	if [ "$force" = true ]; then
		kubectl -n "$ns" delete pod "$new_pod" --ignore-not-found --wait=true || return $?
	fi

	printf '%s\n' "$manifest" | kubectl -n "$ns" create -f - || return $?

	if [ "$exec_after" != true ]; then
		kubectl -n "$ns" get pod "$new_pod"
		return $?
	fi

	if [ "${#exec_cmd[@]}" -eq 0 ]; then
		exec_cmd=(sh)
	fi

	__kcreate_wait_container_running "$ns" "$new_pod" "$container" "$timeout" || return $?

	if [ -n "$container" ]; then
		exec_args=(-c "$container")
	fi

	kubectl -n "$ns" exec -it "$new_pod" "${exec_args[@]}" -- "${exec_cmd[@]}"
	rc=$?

	if [ "$rm_after" = true ]; then
		kubectl -n "$ns" delete pod "$new_pod" --ignore-not-found >/dev/null
	fi

	return "$rc"
}