#!/bin/sh
set -eu

usage() {
  echo "usage: $0 <namespace> <kubeconfig> <context> <session-context>" >&2
  exit 2
}

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

namespace=$1
kubeconfig=$2
context=$3
session_context=$4

remote_runtime_json=$(
  python3 - "$session_context" <<'PY'
import json
import pathlib
import sys

ctx = json.loads(pathlib.Path(sys.argv[1]).read_text())
runtime = (ctx.get("terraform") or {}).get("remote_ec2_runtime") or {}
enabled = (
    (ctx.get("host") or {}).get("mode") == "remote-ec2-k3d"
    and runtime.get("enabled") is True
)
print(json.dumps({
    "enabled": enabled,
    "evidence_dir": str(ctx.get("evidence_dir") or "").strip(),
    "region": str(((ctx.get("target") or {}).get("region")) or "").strip(),
    "role_arn": str(runtime.get("role_arn") or "").strip(),
}))
PY
)

remote_runtime_enabled=$(python3 - "$remote_runtime_json" <<'PY'
import json
import sys
print("true" if json.loads(sys.argv[1])["enabled"] else "false")
PY
)
if [ "$remote_runtime_enabled" = "true" ]; then
  remote_runtime_region=$(python3 - "$remote_runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["region"])
PY
  )
  remote_runtime_evidence_dir=$(python3 - "$remote_runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["evidence_dir"])
PY
  )
  remote_runtime_role_arn=$(python3 - "$remote_runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["role_arn"])
PY
  )
  case "$remote_runtime_role_arn" in
    arn:aws:iam::*:role/*) ;;
    *)
      echo "remote_ec2_runtime.role_arn must be an IAM role ARN" >&2
      exit 1
      ;;
  esac

  KUBECONFIG="$kubeconfig" kubectl --kubeconfig "$kubeconfig" --context "$context" --namespace "$namespace" \
    delete secret aws-creds cp9-aws-credentials --ignore-not-found=true
  for secret_name in aws-creds cp9-aws-credentials; do
    if KUBECONFIG="$kubeconfig" kubectl --kubeconfig "$kubeconfig" --context "$context" --namespace "$namespace" \
      get secret "$secret_name" >/dev/null 2>&1; then
      echo "remote runtime credential Secret still exists: $namespace/$secret_name" >&2
      exit 1
    fi
  done
  if [ -n "$remote_runtime_evidence_dir" ]; then
    python3 - "$remote_runtime_evidence_dir" "$namespace" "$remote_runtime_role_arn" <<'PY'
from datetime import datetime, timezone
import json
import os
import pathlib
import sys
import tempfile

evidence_dir = pathlib.Path(sys.argv[1])
namespace = sys.argv[2]
role_arn = sys.argv[3]
evidence_dir.mkdir(parents=True, exist_ok=True)
payload = {
    "schema_version": 1,
    "captured_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
    "identity_source": "ec2-instance-profile",
    "role_arn": role_arn,
    "target_namespace": namespace,
    "credential_secrets_absent": True,
    "refresh_required": False,
}
target = evidence_dir / f"aws-runtime-identity-{namespace}.json"
with tempfile.NamedTemporaryFile("w", dir=evidence_dir, delete=False) as out:
    json.dump(payload, out, indent=2, sort_keys=True)
    out.write("\n")
    temp_name = out.name
os.chmod(temp_name, 0o600)
os.replace(temp_name, target)
PY
  fi
  echo "remote runtime identity ready: $remote_runtime_role_arn"
  exit 0
fi

creds_file=$(mktemp)
export_stdout=$(mktemp)
export_stderr=$(mktemp)
assume_stdout=$(mktemp)
assume_stderr=$(mktemp)
trap 'rm -f "$creds_file" "$export_stdout" "$export_stderr" "$assume_stdout" "$assume_stderr"' EXIT

runtime_json=$(
  python3 - "$session_context" <<'PY'
import json
import pathlib
import re
import sys

ctx = json.loads(pathlib.Path(sys.argv[1]).read_text())
resolved_path = str(ctx.get("resolved_session_path") or "").strip()
resolved = pathlib.Path(resolved_path).read_text() if resolved_path else ""

def yaml_scalar(key):
    pattern = rf"(?m)^[ \t]*{re.escape(key)}:[ \t]*(.+?)[ \t]*$"
    match = re.search(pattern, resolved)
    if not match:
        return ""
    value = match.group(1).strip()
    if (value.startswith('"') and value.endswith('"')) or (value.startswith("'") and value.endswith("'")):
        value = value[1:-1]
    return value.strip()

profile = yaml_scalar("profile")
source_profile = yaml_scalar("source_profile") or profile
profile = profile or "default"
source_profile = source_profile or profile
role_arn = yaml_scalar("role_arn")
target = ctx.get("target") or {}
session_tags = ctx.get("session_tags") or {}
if not isinstance(session_tags, dict):
    raise SystemExit("CP9_SESSION_CONTEXT.session_tags must be an object")

print(json.dumps({
    "session_name": str(ctx.get("session_name") or "").strip(),
    "owner": str(ctx.get("owner") or "").strip(),
    "cp9_owner_tag": str(session_tags.get("cp9-owner") or "").strip(),
    "cp9_session_tag": str(session_tags.get("cp9-session") or "").strip(),
    "region": str(target.get("region") or "").strip(),
    "profile": profile,
    "source_profile": source_profile,
    "role_arn": role_arn,
}))
PY
)

runtime_profile=$(python3 - "$runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["profile"])
PY
)
source_profile=$(python3 - "$runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["source_profile"])
PY
)
region=$(python3 - "$runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["region"])
PY
)
role_arn=$(python3 - "$runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["role_arn"])
PY
)
duration_seconds=${CP9_AWS_CREDS_DURATION_SECONDS:-43200}
session_name=$(python3 - "$runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["session_name"])
PY
)
owner=$(python3 - "$runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["owner"])
PY
)
cp9_owner_tag=$(python3 - "$runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["cp9_owner_tag"])
PY
)
cp9_session_tag=$(python3 - "$runtime_json" <<'PY'
import json
import sys
print(json.loads(sys.argv[1])["cp9_session_tag"])
PY
)

echo "aws-creds runtime profile: $runtime_profile"
echo "aws-creds source profile: $source_profile"
echo "aws-creds target namespace: $namespace"

set +e
aws configure export-credentials --profile "$runtime_profile" --format env-no-export >"$export_stdout" 2>"$export_stderr"
export_status=$?
set -e
selected_profile=$runtime_profile
assume_role=false
if [ -n "$role_arn" ] && [ "$runtime_profile" = "$source_profile" ]; then
  assume_role=true
fi
if [ "$export_status" -ne 0 ]; then
  if [ "$source_profile" != "$runtime_profile" ]; then
    echo "aws-creds runtime profile export failed; trying source profile: $source_profile" >&2
    if [ -s "$export_stderr" ]; then
      cat "$export_stderr" >&2
    fi
    : >"$export_stdout"
    : >"$export_stderr"
    set +e
    aws configure export-credentials --profile "$source_profile" --format env-no-export >"$export_stdout" 2>"$export_stderr"
    export_status=$?
    set -e
    selected_profile=$source_profile
    assume_role=true
  fi
  if [ "$export_status" -ne 0 ]; then
    if [ -s "$export_stderr" ]; then
      cat "$export_stderr" >&2
    fi
    exit "$export_status"
  fi
fi
echo "aws-creds selected profile: $selected_profile"

access_key_id=""
secret_access_key=""
session_token=""
while IFS='=' read -r key value || [ -n "$key" ]; do
  case "$key" in
    AWS_ACCESS_KEY_ID) access_key_id=$value ;;
    AWS_SECRET_ACCESS_KEY) secret_access_key=$value ;;
    AWS_SESSION_TOKEN) session_token=$value ;;
    "") ;;
  esac
done <"$export_stdout"

if [ -n "$role_arn" ] && [ "$assume_role" = "true" ]; then
  assume_tag_args=""
  if [ -n "$cp9_owner_tag" ] && [ -n "$cp9_session_tag" ]; then
    assume_tag_args="--tags Key=cp9-owner,Value=$cp9_owner_tag Key=cp9-session,Value=$cp9_session_tag --transitive-tag-keys cp9-owner cp9-session"
  elif [ -n "$owner" ] && [ -n "$session_name" ]; then
    assume_tag_args="--tags Key=cp9-owner,Value=$owner Key=cp9-session,Value=$session_name --transitive-tag-keys cp9-owner cp9-session"
  fi
  set +e
  AWS_ACCESS_KEY_ID="$access_key_id" AWS_SECRET_ACCESS_KEY="$secret_access_key" AWS_SESSION_TOKEN="$session_token" AWS_REGION="$region" AWS_DEFAULT_REGION="$region" \
    aws sts assume-role --role-arn "$role_arn" --role-session-name "$session_name" --duration-seconds "$duration_seconds" --output text --query 'Credentials.[AccessKeyId,SecretAccessKey,SessionToken]' $assume_tag_args >"$assume_stdout" 2>"$assume_stderr"
  assume_status=$?
  if [ "$assume_status" -ne 0 ] && [ "$duration_seconds" != "3600" ] && grep -q "exceeds the 1 hour session limit for roles assumed by role chaining" "$assume_stderr"; then
    duration_seconds=3600
    : >"$assume_stdout"
    : >"$assume_stderr"
    AWS_ACCESS_KEY_ID="$access_key_id" AWS_SECRET_ACCESS_KEY="$secret_access_key" AWS_SESSION_TOKEN="$session_token" AWS_REGION="$region" AWS_DEFAULT_REGION="$region" \
      aws sts assume-role --role-arn "$role_arn" --role-session-name "$session_name" --duration-seconds "$duration_seconds" --output text --query 'Credentials.[AccessKeyId,SecretAccessKey,SessionToken]' $assume_tag_args >"$assume_stdout" 2>"$assume_stderr"
    assume_status=$?
  fi
  set -e
  if [ "$assume_status" -ne 0 ]; then
    if [ -s "$assume_stderr" ]; then
      cat "$assume_stderr" >&2
    fi
    exit "$assume_status"
  fi
  set -- $(cat "$assume_stdout")
  access_key_id=${1:-}
  secret_access_key=${2:-}
  session_token=${3:-}
fi

if [ -z "${access_key_id:-}" ] || [ -z "${secret_access_key:-}" ]; then
  echo "selected AWS profile material is incomplete for $source_profile" >&2
  exit 1
fi

{
  printf '[default]\n'
  printf 'aws_access_key_id = %s\n' "$access_key_id"
  printf 'aws_secret_access_key = %s\n' "$secret_access_key"
  if [ -n "${session_token:-}" ]; then
    printf 'aws_session_token = %s\n' "$session_token"
  fi
  if [ -n "$region" ]; then
    printf 'region = %s\n' "$region"
  fi
} >"$creds_file"

KUBECONFIG="$kubeconfig" kubectl --context "$context" -n "$namespace" create secret generic aws-creds --from-file=creds="$creds_file" --dry-run=client -o yaml |
  KUBECONFIG="$kubeconfig" kubectl --context "$context" apply -f -

KUBECONFIG="$kubeconfig" kubectl --context "$context" -n "$namespace" create secret generic cp9-aws-credentials \
  --from-literal=AWS_ACCESS_KEY_ID="$access_key_id" \
  --from-literal=AWS_SECRET_ACCESS_KEY="$secret_access_key" \
  --from-literal=AWS_SESSION_TOKEN="${session_token:-}" \
  --dry-run=client -o yaml |
  KUBECONFIG="$kubeconfig" kubectl --context "$context" apply -f -

KUBECONFIG="$kubeconfig" kubectl --context "$context" -n "$namespace" get secret aws-creds -o name
