#!/bin/sh
set -eu

usage() {
  echo "usage: $0 <kustomize-path> <namespace-or-> <kubeconfig> <context> <old=new>..." >&2
  exit 2
}

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

kustomize_path=$1
namespace=$2
kubeconfig=$3
context=$4
shift 4

tmp=$(mktemp)
new_images=$(mktemp)
trap 'rm -f "$tmp" "$new_images"' EXIT

kubectl kustomize "$kustomize_path" >"$tmp"

for replacement in "$@"; do
  old=${replacement%%=*}
  new=${replacement#*=}
  printf '%s\n' "$new" >>"$new_images"
  python3 - "$tmp" "$old" "$new" <<'PY'
import pathlib
import sys

path = pathlib.Path(sys.argv[1])
old = sys.argv[2]
new = sys.argv[3]
text = path.read_text()
count = text.count(old)
if count == 0:
    print(f"expected image replacement not found: {old}", file=sys.stderr)
    sys.exit(1)
text = text.replace(old, new)
path.write_text(text)
PY
done

if [ "${CP9_PRESERVE_IMAGE_PULL_POLICY:-0}" != "1" ]; then
  python3 - "$tmp" <<'PY'
import pathlib
import sys

path = pathlib.Path(sys.argv[1])
text = path.read_text().replace("imagePullPolicy: Always", "imagePullPolicy: IfNotPresent")
path.write_text(text)
PY
else
  python3 - "$tmp" "$new_images" <<'PY'
import pathlib
import sys

path = pathlib.Path(sys.argv[1])
image_path = pathlib.Path(sys.argv[2])
target_images = {line.strip() for line in image_path.read_text().splitlines() if line.strip()}
lines = path.read_text().splitlines()
out = []
pending_policy_indent = None

def flush_pending():
    global pending_policy_indent
    if pending_policy_indent is not None:
        out.append(" " * pending_policy_indent + "imagePullPolicy: Always")
        pending_policy_indent = None

for line in lines:
    stripped = line.lstrip()
    indent = len(line) - len(stripped)
    if pending_policy_indent is not None:
        if stripped and not stripped.startswith("#"):
            if indent < pending_policy_indent:
                flush_pending()
            elif indent == pending_policy_indent and stripped.startswith("imagePullPolicy:"):
                out.append(" " * pending_policy_indent + "imagePullPolicy: Always")
                pending_policy_indent = None
                continue

    if stripped.startswith("image: "):
        if pending_policy_indent is not None and indent <= pending_policy_indent:
            flush_pending()
        out.append(line)
        image = stripped[len("image: "):].strip().strip('"').strip("'")
        if image in target_images:
            pending_policy_indent = indent
        continue

    out.append(line)

flush_pending()
path.write_text("\n".join(out) + "\n")
PY
fi

if [ -n "${CP9_SESSION_CONTEXT:-}" ] && [ -f "${CP9_SESSION_CONTEXT:-}" ]; then
  python3 - "$tmp" "$CP9_SESSION_CONTEXT" <<'PY'
import base64
import json
import os
import pathlib
import re
import sys
import urllib.parse

manifest_path = pathlib.Path(sys.argv[1])
context_path = pathlib.Path(sys.argv[2])

ctx = json.loads(context_path.read_text())
outputs_path = pathlib.Path(ctx.get("terraform", {}).get("outputs_path", ""))
outputs = {}
if outputs_path.is_file():
    outputs = json.loads(outputs_path.read_text())

def output_value(name, default=""):
    raw = outputs.get(name, {})
    if isinstance(raw, dict) and "value" in raw and raw["value"] is not None:
        return raw["value"]
    return default

session = str(ctx.get("session_name", "")).strip()
namespace = str(ctx.get("namespace", session)).strip()
if len(session) <= 23:
    load_balancer_names = {
        "__CP9_AI_GATEWAY_LB_NAME__": session + "-ai",
        "__CP9_CLUSTER_PROXY_LB_NAME__": session + "-proxy",
        "__CP9_OPS_SERVICE_LB_NAME__": session + "-ops",
        "__CP9_REGIONAL_SERVICE_LB_NAME__": session + "-regional",
        "__CP9_SQL_GATEWAY_LB_NAME__": session + "-sql",
    }
else:
    load_balancer_names = {
        "__CP9_AI_GATEWAY_LB_NAME__": session + "-a",
        "__CP9_CLUSTER_PROXY_LB_NAME__": session + "-p",
        "__CP9_OPS_SERVICE_LB_NAME__": session + "-o",
        "__CP9_REGIONAL_SERVICE_LB_NAME__": session + "-r",
        "__CP9_SQL_GATEWAY_LB_NAME__": session + "-s",
    }
target = ctx.get("target", {})
region = str(target.get("region", "")).strip()
if not region:
    raise SystemExit("Starter remote apply requires target.region")
account = str(target.get("account", "")).strip()
terraform = ctx.get("terraform") or {}
component_pca_arn = str(((terraform.get("session_product_pca") or {}).get("arn")) or "").strip()
runtime_kms_arn = str(((terraform.get("session_cmk") or {}).get("arn")) or "").strip()
remote_eks = output_value("remote_eks", {}) or {}
remote_cluster = str(remote_eks.get("cluster_name") or (session + "-regional")).strip()
runtime_role = str(remote_eks.get("runtime_role_arn") or output_value("iam_role_arn", "")).strip()
runtime_role_name = runtime_role.rsplit("/", 1)[-1] if runtime_role else ""
remote_subnets = remote_eks.get("node_subnet_ids") or []
public_host = urllib.parse.urlparse(
    str((((ctx.get("host") or {}).get("remote") or {}).get("public_https_url")) or "")
).hostname or ""
infra_table = str(output_value("starter_infra_cluster_table", "") or ("cp9-" + session + "-starter-infra-cluster")).strip()
cluster_queue = str(output_value("starter_cluster_queue_url", "") or "").strip()
serverless_storage = output_value("starter_serverless_storage", {}) or {}
serverless_storage_bucket = str(serverless_storage.get("bucket") or ("cp9-" + session + "-starter-cse")).strip()
serverless_storage_prefix = str(serverless_storage.get("prefix") or ("starter/" + session + "/tidb-serverless")).strip()
starter_meta_plane = output_value("starter_meta_plane", {}) or {}
starter_rds_secret_arn = str(starter_meta_plane.get("rds_secret_arn") or "").strip()
starter_infra_rds_secret_arn = str(starter_meta_plane.get("infra_rds_secret_arn") or "").strip()
starter_audit_kms_alias_arn = str(starter_meta_plane.get("audit_kms_alias_arn") or "").strip()
starter_audit_key_ciphertext = str(starter_meta_plane.get("audit_key_ciphertext") or "").strip()
starter_cse_master_key_id = str(starter_meta_plane.get("cse_master_key_id") or "").strip()
starter_cse_key_ciphertext = str(starter_meta_plane.get("cse_key_ciphertext") or "").strip()
starter_cse_key_vendor = str(starter_meta_plane.get("cse_key_vendor") or "").strip()
inference_service_api_key = os.environ.get("CP9_STARTER_INFERENCE_SERVICE_API_KEY", "").strip()
private_inputs = {
    "__CP9_AI_GATEWAY_USERNAME_B64__": "CP9_STARTER_AI_GATEWAY_USERNAME",
    "__CP9_AI_GATEWAY_PASSWORD_B64__": "CP9_STARTER_AI_GATEWAY_PASSWORD",
    "__CP9_FTS_JWT_HOST_B64__": "CP9_STARTER_FTS_JWT_HOST",
    "__CP9_FTS_JWT_AUTH_TOKEN_B64__": "CP9_STARTER_FTS_JWT_AUTH_TOKEN",
    "__CP9_FTS_JWT_ORG_ID_B64__": "CP9_STARTER_FTS_JWT_ORG_ID",
    "__CP9_FTS_JWT_CLUSTER_ID_B64__": "CP9_STARTER_FTS_JWT_CLUSTER_ID",
}
tags = ctx.get("session_tags", {})
owner_tag = str(tags.get("cp9-owner") or ctx.get("owner", "")).strip()
session_tag = str(tags.get("cp9-session") or session).strip()
ttl = str(ctx.get("ttl_expires_at", "")).strip()
if not owner_tag or not session_tag or not ttl:
    raise SystemExit("CP9_SESSION_CONTEXT must provide non-empty cp9-owner, cp9-session, and ttl_expires_at cleanup tags")
lb_tags = ",".join([
    "cp9-owner=" + owner_tag,
    "cp9-session=" + session_tag,
    "cp9-managed=true",
    "cp9-ttl=" + ttl,
])

replacements = {
    "__CP9_SESSION_NAME__": session,
    "__CP9_NAMESPACE__": namespace,
    "__CP9_TARGET_ACCOUNT__": account,
    "__CP9_TARGET_ACCOUNT_ID__": json.dumps(account) if account else "",
    "__CP9_TARGET_REGION__": region,
    "__CP9_REMOTE_CLUSTER_NAME__": remote_cluster,
    "__CP9_REMOTE_EKS_CLUSTER_NAME__": remote_cluster,
    "__CP9_REMOTE_EKS_VPC_ID__": str(remote_eks.get("vpc_id") or "").strip(),
    "__CP9_REMOTE_EKS_SUBNET_ID_0__": str(remote_subnets[0] if len(remote_subnets) > 0 else "").strip(),
    "__CP9_REMOTE_EKS_SUBNET_ID_1__": str(remote_subnets[1] if len(remote_subnets) > 1 else "").strip(),
    "__CP9_REMOTE_EKS_SUBNET_ID_2__": str(remote_subnets[2] if len(remote_subnets) > 2 else "").strip(),
    "__CP9_REMOTE_EKS_ENDPOINT_SECURITY_GROUP_ID__": str(
        remote_eks.get("endpoint_security_group_id") or ""
    ).strip(),
    "__CP9_PUBLIC_HOST__": public_host,
    "__CP9_RUNTIME_ROLE_ARN__": runtime_role,
    "__CP9_RUNTIME_ROLE_NAME__": runtime_role_name,
    "__CP9_INFRA_CLUSTER_TABLE__": infra_table,
    "__CP9_CLUSTER_QUEUE_URL__": cluster_queue,
    "__CP9_SERVERLESS_STORAGE_BUCKET__": serverless_storage_bucket,
    "__CP9_SERVERLESS_STORAGE_PREFIX__": serverless_storage_prefix,
    "__CP9_STARTER_RDS_SECRET_ARN__": starter_rds_secret_arn,
    "__CP9_STARTER_INFRA_RDS_SECRET_ARN__": starter_infra_rds_secret_arn,
    "__CP9_STARTER_AUDIT_KMS_ALIAS_ARN__": starter_audit_kms_alias_arn,
    "__CP9_STARTER_AUDIT_KEY_CIPHERTEXT__": starter_audit_key_ciphertext,
    "CP9_STARTER_CSE_MASTER_KEY_ID": starter_cse_master_key_id,
    "CP9_STARTER_CSE_KEY_CIPHERTEXT": starter_cse_key_ciphertext,
    "CP9_STARTER_CSE_KEY_VENDOR": starter_cse_key_vendor,
    "__CP9_LB_TAGS__": lb_tags,
    "__CP9_OWNER_TAG__": owner_tag,
    "__CP9_SESSION_TAG__": session_tag,
    "__CP9_TTL_TAG__": ttl,
}
replacements.update(load_balancer_names)
if inference_service_api_key:
    replacements["__CP9_INFERENCE_SERVICE_API_KEY_B64__"] = base64.b64encode(
        inference_service_api_key.encode("utf-8")
    ).decode("ascii")
for placeholder, env_name in private_inputs.items():
    value = os.environ.get(env_name, "").strip()
    if value:
        replacements[placeholder] = base64.b64encode(value.encode("utf-8")).decode("ascii")

text = manifest_path.read_text()
o11y_placeholders = (
    "__CP9_SESSION_NAME__",
    "__CP9_TARGET_ACCOUNT_ID__",
    "__CP9_TARGET_REGION__",
    "__CP9_PUBLIC_HOST__",
    "__CP9_RUNTIME_ROLE_ARN__",
    "__CP9_REMOTE_EKS_CLUSTER_NAME__",
    "__CP9_REMOTE_EKS_VPC_ID__",
    "__CP9_REMOTE_EKS_SUBNET_ID_0__",
    "__CP9_REMOTE_EKS_SUBNET_ID_1__",
    "__CP9_REMOTE_EKS_SUBNET_ID_2__",
    "__CP9_REMOTE_EKS_ENDPOINT_SECURITY_GROUP_ID__",
)
missing_o11y = [
    placeholder
    for placeholder in o11y_placeholders
    if placeholder in text and not replacements[placeholder]
]
if missing_o11y:
    raise SystemExit("missing Starter O11Y session values: " + ", ".join(missing_o11y))
if re.search(r"(?m)^\s*pcaArn:\s*.*000000000000.*cp9-session-placeholder\s*$", text):
    if not re.fullmatch(
        rf"arn:aws(?:-[a-z0-9-]+)?:acm-pca:{re.escape(region)}:{account}:"
        r"certificate-authority/[0-9a-fA-F-]+",
        component_pca_arn,
    ):
        raise SystemExit("Starter runtime requires a target-account session_product_pca ARN")
    text = re.sub(
        r"(?m)^(\s*pcaArn:\s*).*$",
        lambda match: match.group(1) + component_pca_arn,
        text,
    )
if re.search(r"(?m)^\s*cmkArn:\s*.*000000000000.*cp9-session-placeholder\s*$", text):
    if not re.fullmatch(
        rf"arn:aws(?:-[a-z0-9-]+)?:kms:{re.escape(region)}:{account}:key/[0-9a-fA-F-]+",
        runtime_kms_arn,
    ):
        raise SystemExit("Starter runtime requires a target-account session_cmk ARN")
    text = re.sub(
        r"(?m)^(\s*cmkArn:\s*).*$",
        lambda match: match.group(1) + runtime_kms_arn,
        text,
    )
oversized_load_balancers = sorted(
    name for placeholder, name in load_balancer_names.items()
    if placeholder in text and len(name) > 32
)
if oversized_load_balancers:
    raise SystemExit(
        "session-prefixed AWS load balancer names exceed the 32-character limit: "
        + ", ".join(oversized_load_balancers)
    )
cse_replacements = {
    "CP9_STARTER_CSE_MASTER_KEY_ID": starter_cse_master_key_id,
    "CP9_STARTER_CSE_KEY_CIPHERTEXT": starter_cse_key_ciphertext,
    "CP9_STARTER_CSE_KEY_VENDOR": starter_cse_key_vendor,
}
missing_cse = sorted(key for key, value in cse_replacements.items() if key in text and not value)
if missing_cse:
    raise SystemExit("Starter CSE master key output missing: " + ", ".join(missing_cse))
for old, new in replacements.items():
    text = text.replace(old, new)
manifest_path.write_text(text)
PY
fi

python3 "$(dirname "$0")/filter-starter-ai-fts.py" \
  "$tmp" "${CP9_STARTER_AI_FTS_ENABLED:-0}"

python3 - "$tmp" <<'PY'
import pathlib
import re
import sys

path = pathlib.Path(sys.argv[1])
text = path.read_text()
unresolved = sorted(set(re.findall(
    r"(?:__CP9_(?:(?:OWNER|SESSION|TTL)_TAG|INFERENCE_SERVICE_API_KEY_B64|AI_GATEWAY_(?:USERNAME|PASSWORD)_B64|FTS_JWT_(?:HOST|AUTH_TOKEN|ORG_ID|CLUSTER_ID)_B64|AI_GATEWAY_(?:CERTIFICATE_ARN|LB_NAME)|(?:CLUSTER_PROXY|OPS_SERVICE|REGIONAL_SERVICE|SQL_GATEWAY)_LB_NAME)__|CP9_STARTER_CSE_(?:MASTER_KEY_ID|KEY_CIPHERTEXT|KEY_VENDOR))",
    text,
)))
if unresolved:
    raise SystemExit("unresolved required runtime placeholders: " + ", ".join(unresolved))
if re.search(
    r"(?m)^\s*(?:pcaArn|cmkArn):\s*.*(?:000000000000|cp9-session-placeholder)\s*$",
    text,
):
    raise SystemExit(
        "Starter runtime manifest retained a PCA/KMS session placeholder; "
        "CP9_SESSION_CONTEXT with target-account Terraform outputs is required"
    )
PY

if [ -n "${CP9_STARTER_PDMS_MODE:-}" ]; then
  python3 "$(dirname "$0")/filter-starter-pdms.py" "$tmp" "$CP9_STARTER_PDMS_MODE"
fi

if [ "${CP9_APPLY_CRDS_ONLY:-0}" = "1" ]; then
  python3 - "$tmp" <<'PY'
import pathlib
import re
import sys

path = pathlib.Path(sys.argv[1])
documents = re.split(r"(?m)^---\s*$", path.read_text())
crds = [
    document.strip()
    for document in documents
    if re.search(r"(?m)^kind:\s*CustomResourceDefinition\s*$", document)
]
if not crds:
    print("expected at least one CustomResourceDefinition in rendered kustomize output", file=sys.stderr)
    sys.exit(1)
path.write_text("\n---\n".join(crds) + "\n")
PY
fi

render_image_pull_secret() {
  image_pull_secret=${CP9_IMAGE_PULL_SECRET:-}
  if [ -z "$image_pull_secret" ] || [ "$namespace" = "-" ] || [ -z "$namespace" ]; then
    return 0
  fi

  python3 - "$tmp" "$image_pull_secret" <<'PY'
import json
import re
import sys
from pathlib import Path

path = Path(sys.argv[1])
secret = sys.argv[2]
text = path.read_text()
pod_spec_paths = {
    "Deployment": ("spec", "template", "spec"),
    "StatefulSet": ("spec", "template", "spec"),
    "DaemonSet": ("spec", "template", "spec"),
    "CronJob": ("spec", "jobTemplate", "spec", "template", "spec"),
}


def add_pull_secret(document: str) -> str:
    kind_match = re.search(r"(?m)^kind:\s*(\S+)\s*$", document)
    if not kind_match or kind_match.group(1) not in pod_spec_paths:
        return document

    lines = document.splitlines()
    cursor = 0
    indent = -2
    for key in pod_spec_paths[kind_match.group(1)]:
        indent += 2
        for index in range(cursor, len(lines)):
            line = lines[index]
            if len(line) - len(line.lstrip()) == indent and line.strip() == key + ":":
                cursor = index + 1
                break
        else:
            raise SystemExit(f"pod spec not found for {kind_match.group(1)}")

    child_indent = indent + 2
    block_end = len(lines)
    image_pull_index = None
    for index in range(cursor, len(lines)):
        line = lines[index]
        if not line.strip():
            continue
        line_indent = len(line) - len(line.lstrip())
        if line_indent <= indent:
            block_end = index
            break
        if line_indent == child_indent and line.strip() == "imagePullSecrets:":
            image_pull_index = index

    entry = " " * child_indent + "- name: " + json.dumps(secret)
    if image_pull_index is None:
        lines[cursor:cursor] = [" " * child_indent + "imagePullSecrets:", entry]
        return "\n".join(lines)

    secret_end = block_end
    for index in range(image_pull_index + 1, block_end):
        line = lines[index]
        if line.strip() and len(line) - len(line.lstrip()) <= child_indent and not line.lstrip().startswith("-"):
            secret_end = index
            break
    if any(re.match(r"\s*-\s*name:\s*[\"']?" + re.escape(secret) + r"[\"']?\s*$", line) for line in lines[image_pull_index + 1 : secret_end]):
        return "\n".join(lines)
    lines.insert(secret_end, entry)
    return "\n".join(lines)


path.write_text("\n---\n".join(add_pull_secret(document) for document in re.split(r"\n---\n", text)))
PY
}

render_image_pull_secret

runtime_config_extra=
runtime_config_extra_kind=
runtime_config_extra_name=
case "$kustomize_path" in
  */starter-regional/cluster-proxy|starter-regional/cluster-proxy)
    runtime_config_extra_kind=Deployment
    runtime_config_extra_name=dbaas-proxy
    runtime_config_extra=$(
      KUBECONFIG="$kubeconfig" kubectl --context "$context" --namespace "$namespace" \
        get secret/proxy-certificate-secret -o json |
        python3 -c '
import json
import sys

item = json.load(sys.stdin)
print(json.dumps({
    "kind": item.get("kind"),
    "name": (item.get("metadata") or {}).get("name"),
    "type": item.get("type"),
    "data": item.get("data") or {},
}, sort_keys=True, separators=(",", ":")))
'
    )
    ;;
  */starter-regional/tidb-serverless|starter-regional/tidb-serverless|\
  */starter-regional/tidb-serverless-pdms-disabled|starter-regional/tidb-serverless-pdms-disabled)
    runtime_config_extra_kind=StatefulSet
    runtime_config_extra_name=tikv-api
    runtime_config_extra=$(
      KUBECONFIG="$kubeconfig" kubectl --context "$context" --namespace "$namespace" \
        get secret/tikv-api-cluster-secret -o json |
        python3 -c '
import json
import sys

item = json.load(sys.stdin)
print(json.dumps({
    "kind": item.get("kind"),
    "name": (item.get("metadata") or {}).get("name"),
    "type": item.get("type"),
    "data": item.get("data") or {},
}, sort_keys=True, separators=(",", ":")))
'
    )
    ;;
esac
runtime_config_stamp=$(
  python3 - "$tmp" <<'PY'
import hashlib
import re
import sys
from pathlib import Path

documents = []
for document in re.split(r"(?m)^---\s*$", Path(sys.argv[1]).read_text()):
    document = document.strip()
    if re.search(r"(?m)^kind:\s*(ConfigMap|Secret)\s*$", document):
        documents.append(document)
if documents:
    print(hashlib.sha256("\n---\n".join(documents).encode()).hexdigest())
PY
)
if [ -n "$runtime_config_stamp" ]; then
  workload_patch=$(
    python3 - "$runtime_config_stamp" <<'PY'
import json
import sys

print(json.dumps({
    "spec": {
        "template": {
            "metadata": {
                "annotations": {
                    "cp9.tidbcloud.com/runtime-config-sha256": sys.argv[1],
                },
            },
        },
    },
}))
PY
  )
  python3 "$(dirname "$0")/../runtime/scripts/patch-local-manifests.py" "$tmp" \
    --kind Deployment --kind StatefulSet --kind DaemonSet \
    --patch "$workload_patch"
fi
if [ -n "$runtime_config_extra" ]; then
  external_config_stamp=$(
    python3 - "$runtime_config_extra" <<'PY'
import hashlib
import sys

print(hashlib.sha256(sys.argv[1].encode()).hexdigest())
PY
  )
  external_patch=$(
    python3 - "$external_config_stamp" <<'PY'
import json
import sys

print(json.dumps({
    "spec": {
        "template": {
            "metadata": {
                "annotations": {
                    "cp9.tidbcloud.com/external-config-sha256": sys.argv[1],
                },
            },
        },
    },
}))
PY
  )
  python3 "$(dirname "$0")/../runtime/scripts/patch-local-manifests.py" "$tmp" \
    --kind "$runtime_config_extra_kind" \
    --name "$runtime_config_extra_name" \
    --patch "$external_patch" \
    --require-match
fi

if [ "$namespace" = "-" ] || [ -z "$namespace" ]; then
  KUBECONFIG="$kubeconfig" kubectl --context "$context" apply --server-side --force-conflicts -f "$tmp"
else
  KUBECONFIG="$kubeconfig" kubectl --context "$context" --namespace "$namespace" apply --server-side --force-conflicts -f "$tmp"
fi
