#!/usr/bin/env bash
# Pacific AI — AWS Marketplace One-Line Installer
#
# Usage:
#   curl -sSL https://<host>/install.sh | bash
#   bash install.sh                         # run directly
#   bash install.sh --dry-run               # preview steps without executing
#
# Requirements: aws-cli, kubectl, helm >= 3.7.1, eksctl

set -euo pipefail

# ─────────────────────────────────────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────────────────────────────────────
readonly MARKETPLACE_REGISTRY="709825985650.dkr.ecr.us-east-1.amazonaws.com"
readonly CHART_REPO="oci://${MARKETPLACE_REGISTRY}/pacific-ai/pacific-ai"
readonly CHART_WORK_DIR="awsmp-chart"
readonly CREDENTIALS_FILE="./pacific-ai-credentials.txt"
readonly HELM_MIN_VERSION="3.7.1"
readonly DEFAULT_NAMESPACE="pacific-ai"
readonly DEFAULT_CLUSTER_NAME="pacific-ai"
readonly DEFAULT_REGION="us-east-1"
readonly TOTAL_STEPS=11

# ANSI colors
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
BLUE='\033[0;34m'; CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'

# Global state
DRY_RUN=false
TRACK_TIME=false
START_TIME=0
STEP=0
CLUSTER_NAME=""
NAMESPACE=""
REGION=""
BASE_URL=""
JUDGE_PROVIDER=""; JUDGE_API_KEY=""; JUDGE_ENDPOINT=""; JUDGE_MODEL_NAME=""
GENERAL_PROVIDER=""; GENERAL_API_KEY=""; GENERAL_ENDPOINT=""; GENERAL_MODEL_NAME=""
CHART_VERSION=""

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────
info()    { echo -e "${CYAN}ℹ  $*${NC}"; }
success() { echo -e "${GREEN}✓  $*${NC}"; }
warn()    { echo -e "${YELLOW}⚠  $*${NC}"; }
error()   { echo -e "${RED}✗  Error: $*${NC}" >&2; }

step() {
  STEP=$((STEP + 1))
  echo ""
  echo -e "${BOLD}${BLUE}── [${STEP}/${TOTAL_STEPS}] $* ──────────────────────────────────────────${NC}"
}

on_error() {
  echo ""
  error "Installation failed (line $1)."
  echo -e "${YELLOW}Check the output above for details, then re-run this script.${NC}"
  exit 1
}
trap 'on_error $LINENO' ERR

# Read interactively even when the script is piped (curl | bash uses stdin for the pipe)
# Uses printf -v (bash 3.1+) for safe variable-by-name assignment — avoids eval quoting issues.
prompt() {
  local _varname="$1"
  local message="$2"
  local default="${3:-}"
  if [[ -n "$default" ]]; then
    printf "${YELLOW}  %s [%s]: ${NC}" "$message" "$default" >/dev/tty
  else
    printf "${YELLOW}  %s: ${NC}" "$message" >/dev/tty
  fi
  local _val=""
  IFS= read -r -e _val </dev/tty || true
  _val="${_val%$'\r'}"   # strip carriage return (CR+LF terminals)
  # Strip any ANSI escape sequences (e.g. from arrow keys on non-readline terminals)
  _val=$(printf '%s' "$_val" | sed 's/\x1b\[[0-9;]*[a-zA-Z]//g')
  _val="${_val:-$default}"
  printf -v "$_varname" '%s' "$_val"
}

prompt_secret() {
  local _varname="$1"
  local message="$2"
  printf "${YELLOW}  %s: ${NC}" "$message" >/dev/tty
  local _val=""
  IFS= read -rs _val </dev/tty || true
  echo >/dev/tty
  _val="${_val%$'\r'}"
  printf -v "$_varname" '%s' "$_val"
}

# prompt_menu VAR "Question" "Option 1" "Option 2" ...
# Sets VAR to 0-based index of the selection
prompt_menu() {
  local _varname="$1"
  local question="$2"
  shift 2
  local options=("$@")
  echo -e "${YELLOW}  ${question}${NC}" >/dev/tty
  for i in "${!options[@]}"; do
    printf "    %d) %s\n" "$((i + 1))" "${options[$i]}" >/dev/tty
  done
  local _choice=""
  while true; do
    printf "${YELLOW}  Enter number [1]: ${NC}" >/dev/tty
    IFS= read -r _choice </dev/tty || true
    _choice="${_choice%$'\r'}"
    _choice="${_choice:-1}"
    if [[ "$_choice" =~ ^[0-9]+$ ]] && (( _choice >= 1 && _choice <= ${#options[@]} )); then
      break
    fi
    echo -e "${RED}  Invalid choice, please enter a number between 1 and ${#options[@]}${NC}" >/dev/tty
  done
  printf -v "$_varname" '%d' "$(( _choice - 1 ))"
}

prompt_confirm() {
  # Returns 0 (true) on yes, 1 (false) on no
  local message="$1"
  local default="${2:-yes}"
  local _confirm_val=""
  prompt _confirm_val "$message" "$default"
  [[ "$_confirm_val" =~ ^[Yy] ]]
}

gen_password()  { openssl rand -hex 16; }
gen_uuid_like() {
  printf '%s-%s-%s-%s-%s' \
    "$(openssl rand -hex 4)" "$(openssl rand -hex 2)" \
    "$(openssl rand -hex 2)" "$(openssl rand -hex 2)" \
    "$(openssl rand -hex 6)"
}

# Returns 0 if $1 >= $2 (semver)
version_ge() { printf '%s\n' "$2" "$1" | sort -V -C; }

run() {
  if [[ "$DRY_RUN" == "true" ]]; then
    echo -e "${CYAN}  [dry-run] $*${NC}"
  else
    "$@"
  fi
}

# ─────────────────────────────────────────────────────────────────────────────
# Parse arguments
# ─────────────────────────────────────────────────────────────────────────────
parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dry-run) DRY_RUN=true; warn "Dry-run mode: no changes will be made." ;;
      --time)    TRACK_TIME=true ;;
      -h|--help)
        cat <<'EOF'
Pacific AI AWS Marketplace Installer

Usage:
  bash install.sh [OPTIONS]

Options:
  --dry-run   Preview all steps without executing any commands
  -h, --help  Show this help message
EOF
        exit 0 ;;
      *) error "Unknown argument: $1"; exit 1 ;;
    esac
    shift
  done
}

# ─────────────────────────────────────────────────────────────────────────────
# Banner
# ─────────────────────────────────────────────────────────────────────────────
print_banner() {
  echo -e "${BOLD}${BLUE}"
  cat <<'BANNER'
╔═══════════════════════════════════════════════════════════════════╗
║                                                                   ║
║           Pacific AI — AWS Marketplace Installer                  ║
║                                                                   ║
╚═══════════════════════════════════════════════════════════════════╝
BANNER
  echo -e "${NC}"
  echo -e "This script will guide you through deploying Pacific AI on AWS EKS."
  echo -e "Estimated time: ${BOLD}35–55 min${NC} (cluster creation adds ~20 min)\n"
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 1: Prerequisites
# ─────────────────────────────────────────────────────────────────────────────
check_prerequisites() {
  step "Checking prerequisites"

  local missing=()

  # aws CLI
  if ! command -v aws &>/dev/null; then
    missing+=("aws-cli")
    error "AWS CLI not found."
    echo "    Install: https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html"
  else
    success "aws CLI: $(aws --version 2>&1 | head -1)"
  fi

  # kubectl
  if ! command -v kubectl &>/dev/null; then
    missing+=("kubectl")
    error "kubectl not found."
    echo "    Install: https://kubernetes.io/docs/tasks/tools/"
  else
    local _kv
    _kv=$(kubectl version --client --short 2>/dev/null || kubectl version --client 2>/dev/null | head -1 || echo "unknown")
    success "kubectl: ${_kv}"
  fi

  # helm
  if ! command -v helm &>/dev/null; then
    missing+=("helm")
    error "Helm not found."
    echo "    Install: https://helm.sh/docs/intro/install/"
  else
    local helm_ver
    helm_ver=$(helm version --short 2>/dev/null | sed 's/v\([0-9][0-9.]*\).*/\1/')
    if version_ge "$helm_ver" "$HELM_MIN_VERSION"; then
      success "helm: v${helm_ver}"
    else
      missing+=("helm >= ${HELM_MIN_VERSION}")
      error "Helm ${helm_ver} found but ${HELM_MIN_VERSION}+ is required."
      echo "    Upgrade: https://helm.sh/docs/intro/install/"
    fi
  fi

  # eksctl
  if ! command -v eksctl &>/dev/null; then
    missing+=("eksctl")
    error "eksctl not found."
    echo "    Install: https://eksctl.io/installation/"
  else
    success "eksctl: $(eksctl version 2>/dev/null)"
  fi

  if [[ ${#missing[@]} -gt 0 ]]; then
    echo ""
    error "Missing required tools: ${missing[*]}"
    echo "  Please install them and re-run this script."
    exit 1
  fi

  success "All prerequisites satisfied."
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 2: AWS Credentials
# ─────────────────────────────────────────────────────────────────────────────
check_aws_credentials() {
  step "Verifying AWS credentials"

  local identity
  if ! identity=$(aws sts get-caller-identity 2>&1); then
    error "AWS credentials are not configured or are invalid."
    echo ""
    echo "  Configure credentials using one of:"
    echo "    • aws configure"
    echo "    • Export AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY (and AWS_SESSION_TOKEN if using STS)"
    echo "    • Attach an IAM role to your EC2 / Cloud9 instance"
    exit 1
  fi

  local account_id user_arn
  account_id=$(echo "$identity" | grep -o '"Account": "[^"]*"' | cut -d'"' -f4)
  user_arn=$(echo "$identity"   | grep -o '"Arn": "[^"]*"'     | cut -d'"' -f4)

  success "AWS Account:  ${account_id}"
  success "IAM Identity: ${user_arn}"
  echo ""
  warn "Please confirm this is the correct AWS account before continuing."
  if ! prompt_confirm "Continue with this account? (yes/no)" "yes"; then
    echo "  Aborted. Configure the correct AWS credentials and re-run."
    exit 0
  fi
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 3: Region
# ─────────────────────────────────────────────────────────────────────────────
collect_region() {
  step "AWS Region"
  local default_region="${AWS_DEFAULT_REGION:-${AWS_REGION:-$DEFAULT_REGION}}"
  prompt REGION "AWS region to deploy to" "$default_region"
  export AWS_DEFAULT_REGION="$REGION"
  success "Region: ${REGION}"
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 4: EKS Cluster
# ─────────────────────────────────────────────────────────────────────────────
setup_cluster() {
  step "EKS Cluster"

  local cluster_choice
  prompt_menu cluster_choice "How would you like to proceed?" \
    "Create a new EKS cluster" \
    "Use an existing EKS cluster"

  if [[ "$cluster_choice" == "0" ]]; then
    _create_new_cluster
  else
    _connect_existing_cluster
  fi
}

_connect_existing_cluster() {
  echo ""
  prompt CLUSTER_NAME "Enter your EKS cluster name"
  if [[ -z "$CLUSTER_NAME" ]]; then
    error "Cluster name cannot be empty."
    exit 1
  fi

  info "Updating kubeconfig for cluster '${CLUSTER_NAME}'..."
  run aws eks update-kubeconfig --name "$CLUSTER_NAME" --region "$REGION"

  # OIDC is required for IRSA (IAM Roles for Service Accounts), which is used by
  # the EFS CSI driver and the Marketplace service account. This is a no-op if
  # OIDC is already associated with the cluster.
  info "Ensuring OIDC provider is associated (required for IAM service accounts)..."
  run eksctl utils associate-iam-oidc-provider \
    --cluster "$CLUSTER_NAME" \
    --region "$REGION" \
    --approve

  echo ""
  info "Cluster nodes:"
  run kubectl get nodes
  success "Connected to cluster '${CLUSTER_NAME}'."
}

_create_new_cluster() {
  echo ""
  echo -e "${BOLD}  New EKS Cluster Configuration${NC}"
  echo ""

  prompt CLUSTER_NAME "Cluster name" "$DEFAULT_CLUSTER_NAME"

  echo ""
  echo "  Recommended instance types:"
  echo "    1) t3.2xlarge   — 8 vCPU, 32 GB RAM                      (cost-effective)"
  echo "    2) m5.2xlarge   — 8 vCPU, 32 GB RAM                      (general purpose)"
  echo "    3) g4dn.xlarge  — 4 vCPU, 16 GB RAM + 1× NVIDIA T4 GPU   (LLM inference)"
  echo "    4) Custom"
  echo ""

  local instance_choice
  prompt_menu instance_choice "Select instance type" \
    "t3.2xlarge  (CPU, cost-effective)" \
    "m5.2xlarge  (CPU, general purpose)" \
    "g4dn.xlarge (GPU, recommended for LLM inference)" \
    "Custom instance type"

  local node_type
  case "$instance_choice" in
    0) node_type="t3.2xlarge"  ;;
    1) node_type="m5.2xlarge"  ;;
    2) node_type="g4dn.xlarge" ;;
    3) prompt node_type "Instance type (e.g. m5.4xlarge)" "" ;;
  esac

  local node_count k8s_version
  prompt node_count "Number of worker nodes" "2"

  local valid_k8s_versions=("1.33" "1.34" "1.35")
  while true; do
    prompt k8s_version "Kubernetes version (1.33 / 1.34 / 1.35)" "1.35"
    local _valid=false
    for _v in "${valid_k8s_versions[@]}"; do
      [[ "$k8s_version" == "$_v" ]] && _valid=true && break
    done
    if [[ "$_valid" == "true" ]]; then
      break
    fi
    echo -e "${RED}  Invalid version '${k8s_version}'. Must be one of: 1.33, 1.34, 1.35${NC}" >/dev/tty
  done

  echo ""
  echo -e "  ${BOLD}Cluster summary:${NC}"
  echo "    Name:           ${CLUSTER_NAME}"
  echo "    Region:         ${REGION}"
  echo "    Instance type:  ${node_type}"
  echo "    Node count:     ${node_count}"
  echo "    K8s version:    ${k8s_version}"
  echo ""
  warn "Creating an EKS cluster takes approximately 15 minutes."
  if ! prompt_confirm "Proceed with cluster creation? (yes/no)" "yes"; then
    echo "  Aborted."
    exit 0
  fi

  # Generate an eksctl config file so we can explicitly set autoModeConfig.enabled: false.
  # This silences the "Auto Mode will be enabled by default" warning and future-proofs the
  # install against upcoming eksctl releases that would change this default.
  local cluster_config_file
  cluster_config_file=$(mktemp /tmp/eksctl-pacific-ai-XXXXXX.yaml)
  cat > "$cluster_config_file" <<EKSCTL_CONFIG
apiVersion: eksctl.io/v1alpha5
kind: ClusterConfig

metadata:
  name: ${CLUSTER_NAME}
  region: ${REGION}
  version: "${k8s_version}"

iam:
  withOIDC: true

autoModeConfig:
  enabled: false

managedNodeGroups:
  - name: pacific-ai-nodes
    instanceType: ${node_type}
    desiredCapacity: ${node_count}
    minSize: 1
    maxSize: $((node_count * 2))
EKSCTL_CONFIG

  info "Creating EKS cluster '${CLUSTER_NAME}'..."
  run eksctl create cluster -f "$cluster_config_file"
  rm -f "$cluster_config_file"

  echo ""
  info "Cluster nodes:"
  run kubectl get nodes
  success "Cluster '${CLUSTER_NAME}' created and configured."
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 5: Namespace
# ─────────────────────────────────────────────────────────────────────────────
setup_namespace() {
  step "Kubernetes namespace"
  prompt NAMESPACE "Namespace to deploy Pacific AI into" "$DEFAULT_NAMESPACE"
  run kubectl create namespace "$NAMESPACE" --dry-run=client -o yaml | run kubectl apply -f -
  run kubectl config set-context --current --namespace="$NAMESPACE"
  success "Namespace '${NAMESPACE}' is ready and set as the default for this context."
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 6: EFS Storage
# ─────────────────────────────────────────────────────────────────────────────
setup_storage() {
  step "Storage (Amazon EFS)"

  if kubectl get storageclass pacific-ai-efs &>/dev/null; then
    success "StorageClass 'pacific-ai-efs' already exists — skipping storage setup."
    return
  fi

  info "Setting up Amazon EFS for persistent storage."
  info "Resources created: Security Group, EFS Filesystem, Mount Targets, EFS CSI Driver."
  echo ""

  # ── Gather cluster network info ────────────────────────────────────────────
  info "Fetching cluster network information..."
  local vpc_id subnet_ids vpc_cidr
  vpc_id=$(aws eks describe-cluster \
    --name "$CLUSTER_NAME" --region "$REGION" \
    --query 'cluster.resourcesVpcConfig.vpcId' --output text)
  subnet_ids=$(aws eks describe-cluster \
    --name "$CLUSTER_NAME" --region "$REGION" \
    --query 'cluster.resourcesVpcConfig.subnetIds[]' --output text)
  vpc_cidr=$(aws ec2 describe-vpcs \
    --vpc-ids "$vpc_id" --region "$REGION" \
    --query 'Vpcs[0].CidrBlock' --output text)
  info "VPC:     ${vpc_id}"
  info "CIDR:    ${vpc_cidr}"

  # ── EFS Security Group ─────────────────────────────────────────────────────
  info "Creating EFS security group (NFS port 2049 from VPC CIDR)..."
  local efs_sg_id
  if [[ "$DRY_RUN" == "true" ]]; then
    echo -e "   ${YELLOW}[dry-run]${NC} aws ec2 create-security-group --group-name ${CLUSTER_NAME}-efs-sg"
    efs_sg_id="sg-dryrun00000000"
  else
    efs_sg_id=$(aws ec2 create-security-group \
      --group-name "${CLUSTER_NAME}-efs-sg" \
      --description "EFS security group for Pacific AI cluster ${CLUSTER_NAME}" \
      --vpc-id "$vpc_id" \
      --region "$REGION" \
      --query 'GroupId' --output text)
    aws ec2 authorize-security-group-ingress \
      --group-id "$efs_sg_id" \
      --protocol tcp --port 2049 \
      --cidr "$vpc_cidr" \
      --region "$REGION"
    aws ec2 create-tags \
      --resources "$efs_sg_id" \
      --tags "Key=Name,Value=${CLUSTER_NAME}-efs-sg" \
      --region "$REGION"
  fi
  success "EFS security group: ${efs_sg_id}"

  # ── EFS Filesystem ─────────────────────────────────────────────────────────
  info "Creating EFS filesystem (encrypted)..."
  local efs_id
  if [[ "$DRY_RUN" == "true" ]]; then
    echo -e "   ${YELLOW}[dry-run]${NC} aws efs create-file-system"
    efs_id="fs-dryrun00000000000"
  else
    efs_id=$(aws efs create-file-system \
      --performance-mode generalPurpose \
      --throughput-mode bursting \
      --encrypted \
      --region "$REGION" \
      --tags "Key=Name,Value=${CLUSTER_NAME}-efs" \
      --query 'FileSystemId' --output text)
  fi
  success "EFS filesystem: ${efs_id}"

  # Wait for the filesystem to be available before creating mount targets
  if [[ "$DRY_RUN" != "true" ]]; then
    info "Waiting for EFS filesystem to become available..."
    local retries=30
    for ((i = 1; i <= retries; i++)); do
      local state
      state=$(aws efs describe-file-systems \
        --file-system-id "$efs_id" --region "$REGION" \
        --query 'FileSystems[0].LifeCycleState' --output text)
      if [[ "$state" == "available" ]]; then
        echo ""
        success "EFS filesystem is available."
        break
      fi
      printf "  [%d/%d] State: %s — waiting...\r" "$i" "$retries" "$state"
      sleep 5
    done
  fi

  # ── Mount Targets (one per subnet) ─────────────────────────────────────────
  info "Creating EFS mount targets in cluster subnets..."
  local subnet
  for subnet in $subnet_ids; do
    if [[ "$DRY_RUN" == "true" ]]; then
      echo -e "   ${YELLOW}[dry-run]${NC} aws efs create-mount-target --subnet-id ${subnet}"
    else
      aws efs create-mount-target \
        --file-system-id "$efs_id" \
        --subnet-id "$subnet" \
        --security-groups "$efs_sg_id" \
        --region "$REGION" &>/dev/null \
        || warn "Mount target in ${subnet} may already exist — skipping."
    fi
  done
  success "Mount targets created."

  # ── EBS CSI Driver (required for MongoDB gp2 volumes) ──────────────────────
  info "Installing AWS EBS CSI driver add-on (required for MongoDB storage)..."
  if [[ "$DRY_RUN" != "true" ]]; then
    run eksctl create iamserviceaccount \
      --name ebs-csi-controller-sa \
      --namespace kube-system \
      --cluster "$CLUSTER_NAME" \
      --region "$REGION" \
      --attach-policy-arn arn:aws:iam::aws:policy/service-role/AmazonEBSCSIDriverPolicy \
      --approve \
      --override-existing-serviceaccounts

    local ebs_role_arn
    ebs_role_arn=$(kubectl get sa ebs-csi-controller-sa -n kube-system \
      -o jsonpath='{.metadata.annotations.eks\.amazonaws\.com/role-arn}')

    run aws eks create-addon \
      --cluster-name "$CLUSTER_NAME" \
      --region "$REGION" \
      --addon-name aws-ebs-csi-driver \
      --service-account-role-arn "$ebs_role_arn" \
      --resolve-conflicts OVERWRITE

    info "Waiting for EBS CSI driver to become active..."
    run aws eks wait addon-active \
      --cluster-name "$CLUSTER_NAME" \
      --region "$REGION" \
      --addon-name aws-ebs-csi-driver
  else
    echo -e "   ${YELLOW}[dry-run]${NC} eksctl create iamserviceaccount ebs-csi-controller-sa"
    echo -e "   ${YELLOW}[dry-run]${NC} aws eks create-addon aws-ebs-csi-driver"
  fi
  success "EBS CSI driver is ready."

  # ── EFS CSI Driver ─────────────────────────────────────────────────────────
  info "Installing AWS EFS CSI driver add-on..."
  run eksctl create addon \
    --name aws-efs-csi-driver \
    --cluster "$CLUSTER_NAME" \
    --region "$REGION" \
    --force

  info "Waiting for EFS CSI controller to be ready..."
  if [[ "$DRY_RUN" != "true" ]]; then
    local retries=30
    for ((i = 1; i <= retries; i++)); do
      if kubectl get deployment efs-csi-controller -n kube-system &>/dev/null; then
        break
      fi
      printf "  [%d/%d] Waiting for deployment to appear...\r" "$i" "$retries"
      sleep 5
    done
    echo ""
  fi
  run kubectl rollout status deployment efs-csi-controller -n kube-system --timeout=120s
  success "EFS CSI driver is ready."

  # ── StorageClass ───────────────────────────────────────────────────────────
  info "Applying 'pacific-ai-efs' StorageClass as cluster default..."
  local sc_manifest
  sc_manifest=$(mktemp /tmp/pacific-ai-efs-sc-XXXXXX.yaml)
  cat > "$sc_manifest" <<EOF
apiVersion: storage.k8s.io/v1
kind: StorageClass
metadata:
  name: pacific-ai-efs
  annotations:
    storageclass.kubernetes.io/is-default-class: "true"
parameters:
  provisioningMode: efs-ap
  fileSystemId: "${efs_id}"
  directoryPerms: "777"
  uid: "0"
  gid: "0"
  basePath: "/pacific-ai"
provisioner: efs.csi.aws.com
reclaimPolicy: Retain
allowVolumeExpansion: true
volumeBindingMode: Immediate
EOF
  run kubectl apply -f "$sc_manifest"
  rm -f "$sc_manifest"
  success "StorageClass 'pacific-ai-efs' is now the cluster default."
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 7: IAM Service Account (AWS Marketplace requirement)
# ─────────────────────────────────────────────────────────────────────────────
attach_iam_to_service_account() {
  step "IAM Service Account (AWS Marketplace)"

  # Helm has already created the 'pacific-ai-pacific-ai' service account as part of the
  # Helm release. We now attach the required AWS Marketplace IAM policies to it using
  # eksctl with --override-existing-serviceaccounts, which adds the IAM role ARN
  # annotation to the existing SA without changing its other Helm-managed properties.
  info "Attaching AWS Marketplace IAM policies to the Helm-managed service account..."
  echo ""

  run eksctl create iamserviceaccount \
    --name pacific-ai-pacific-ai \
    --namespace "$NAMESPACE" \
    --cluster "$CLUSTER_NAME" \
    --region "$REGION" \
    --attach-policy-arn arn:aws:iam::aws:policy/AWSMarketplaceMeteringFullAccess \
    --attach-policy-arn arn:aws:iam::aws:policy/AWSMarketplaceMeteringRegisterUsage \
    --attach-policy-arn arn:aws:iam::aws:policy/service-role/AWSLicenseManagerConsumptionPolicy \
    --approve \
    --override-existing-serviceaccounts

  # Restart pods that use this SA so they pick up the new IAM role annotation immediately.
  info "Restarting application pods to apply the new IAM role..."
  run kubectl rollout restart deployment -n "$NAMESPACE" 2>/dev/null || true

  success "IAM role attached to service account 'pacific-ai-pacific-ai'."
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 7: ECR Login + Pull Helm Chart
# ─────────────────────────────────────────────────────────────────────────────
pull_helm_chart() {
  step "Pulling Helm chart from AWS Marketplace ECR"

  info "Authenticating to Marketplace ECR registry..."
  run bash -c "aws ecr get-login-password --region us-east-1 \
    | helm registry login --username AWS --password-stdin '${MARKETPLACE_REGISTRY}'"
  success "Helm registry login successful."

  echo ""
  if [[ -z "$CHART_VERSION" ]]; then
    info "Fetching latest chart version from Pacific AI..."
    CHART_VERSION=$(curl -fsSL "https://marketplace-pacific-ai.s3.us-east-1.amazonaws.com/VERSION" 2>/dev/null || true)
    if [[ -z "$CHART_VERSION" ]]; then
      prompt CHART_VERSION "Pacific AI chart version to install (e.g. 1.0.0)" ""
    else
      info "Latest version: ${CHART_VERSION}"
      prompt CHART_VERSION "Pacific AI chart version to install" "$CHART_VERSION"
    fi
  fi
  if [[ -z "$CHART_VERSION" ]]; then
    error "Chart version is required. Check your AWS Marketplace subscription for the available version."
    exit 1
  fi

  info "Pulling Helm chart v${CHART_VERSION}..."
  rm -rf "${CHART_WORK_DIR}"
  mkdir -p "${CHART_WORK_DIR}"
  run helm pull "${CHART_REPO}" --version "${CHART_VERSION}" --destination "${CHART_WORK_DIR}"
  run bash -c "cd '${CHART_WORK_DIR}' && tar xf ./*.tgz && find . -maxdepth 1 -type f -name '*.tgz' -delete"

  success "Helm chart extracted to ./${CHART_WORK_DIR}/."
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 8: Application Configuration
# ─────────────────────────────────────────────────────────────────────────────
collect_configuration() {
  step "Application configuration"

  # ── Base URL ────────────────────────────────────────────────────────────────
  echo -e "  ${BOLD}Base URL${NC}"
  echo "  This is the URL your users will use to access Pacific AI."
  echo "  Leave blank to auto-detect the LoadBalancer URL after installation."
  echo ""
  prompt BASE_URL "Base URL (e.g. https://pacific-ai.mycompany.com or leave blank)" ""
  echo ""

}

# ─────────────────────────────────────────────────────────────────────────────
# Step 9: Helm Install
# ─────────────────────────────────────────────────────────────────────────────
install_pacific_ai() {
  step "Installing Pacific AI"

  # Generate all secrets
  info "Generating secure random passwords..."
  local kc_admin_pw;          kc_admin_pw=$(gen_password)
  local kc_user_admin;        kc_user_admin=$(gen_password)
  local kc_user_member;       kc_user_member=$(gen_password)
  local kc_user_gov;          kc_user_gov=$(gen_password)
  local kc_user_policy;       kc_user_policy=$(gen_password)
  local kc_user_vendor;       kc_user_vendor=$(gen_password)
  local kc_user_compliance;   kc_user_compliance=$(gen_password)
  local kc_user_risk;         kc_user_risk=$(gen_password)
  local kc_user_developer;    kc_user_developer=$(gen_password)
  local kc_client_sec;        kc_client_sec=$(gen_uuid_like)
  local kc_pg_pw;        kc_pg_pw=$(gen_password)
  local kc_pg_root_pw;   kc_pg_root_pw=$(gen_password)
  local pg_pw;           pg_pw=$(gen_password)
  local redis_pw;        redis_pw=$(gen_password)
  local mongo_root_pw;   mongo_root_pw=$(gen_password)
  local mongo_pw;        mongo_pw=$(gen_password)
  local llm_gw_key;      llm_gw_key="sk-$(gen_password)"
  local llm_gw_redis_pw; llm_gw_redis_pw=$(gen_password)

  local helm_base_url="${BASE_URL:-http://pending}"
  local chart_dir
  chart_dir="$(ls -d "${CHART_WORK_DIR}"/*/)"

  info "Running Helm install (timeout: 15 min)..."
  echo ""

  run helm install pacific-ai "${chart_dir}" \
    --namespace "$NAMESPACE" \
    --create-namespace \
    --timeout 15m \
    --set "global.baseUrl=${helm_base_url}" \
    --set "keycloak.auth.adminPassword=${kc_admin_pw}" \
    --set "keycloakconfig.config.users.passwords.admin=${kc_user_admin}" \
    --set "keycloakconfig.config.users.passwords.member=${kc_user_member}" \
    --set "keycloakconfig.config.users.passwords.governance_officer=${kc_user_gov}" \
    --set "keycloakconfig.config.users.passwords.policy_manager=${kc_user_policy}" \
    --set "keycloakconfig.config.users.passwords.vendor_manager=${kc_user_vendor}" \
    --set "keycloakconfig.config.users.passwords.compliance_officer=${kc_user_compliance}" \
    --set "keycloakconfig.config.users.passwords.risk_manager=${kc_user_risk}" \
    --set "keycloakconfig.config.users.passwords.developer=${kc_user_developer}" \
    --set "keycloak.config.clientSecret=${kc_client_sec}" \
    --set "keycloak.postgresql.auth.password=${kc_pg_pw}" \
    --set "keycloak.postgresql.auth.postgresPassword=${kc_pg_root_pw}" \
    --set "postgresql.auth.password=${pg_pw}" \
    --set "redis.auth.existingSecretPassword=${redis_pw}" \
    --set "mongodb.auth.rootPassword=${mongo_root_pw}" \
    --set "mongodb.auth.password=${mongo_pw}" \
    --set "mongodb.persistence.storageClass=gp2" \
    --set "llm-gateway.masterkey=${llm_gw_key}" \
    --set "llm-gateway.redis.auth.existingSecretPassword=${llm_gw_redis_pw}" \
    --set "pacific.configuration.judge.llm.modelProvider=${JUDGE_PROVIDER}" \
    --set "pacific.configuration.judge.llm.apiKey=${JUDGE_API_KEY}" \
    --set "pacific.configuration.judge.llm.endpoint=${JUDGE_ENDPOINT}" \
    --set "pacific.configuration.judge.llm.modelName=${JUDGE_MODEL_NAME}" \
    --set "pacific.configuration.general.llm.modelProvider=${GENERAL_PROVIDER}" \
    --set "pacific.configuration.general.llm.apiKey=${GENERAL_API_KEY}" \
    --set "pacific.configuration.general.llm.endpoint=${GENERAL_ENDPOINT}" \
    --set "pacific.configuration.general.llm.modelName=${GENERAL_MODEL_NAME}"

  success "Helm install completed."

  # Save credentials immediately so they are not lost even if URL detection fails
  _save_credentials "$kc_admin_pw" \
    "$kc_user_admin" "$kc_user_member" "$kc_user_gov" "$kc_user_policy" \
    "$kc_user_vendor" "$kc_user_compliance" "$kc_user_risk" "$kc_user_developer"

  # Auto-detect LoadBalancer URL if none was provided
  if [[ -z "$BASE_URL" ]]; then
    _detect_loadbalancer_url "${chart_dir}" \
      "$kc_admin_pw" \
      "$kc_user_admin" "$kc_user_member" "$kc_user_gov" "$kc_user_policy" \
      "$kc_user_vendor" "$kc_user_compliance" "$kc_user_risk" "$kc_user_developer"
  fi
}

_detect_loadbalancer_url() {
  local chart_dir="$1"; shift
  local kc_admin_pw="$1"
  local kc_user_admin="$2" kc_user_member="$3" kc_user_gov="$4" kc_user_policy="$5"
  local kc_user_vendor="$6" kc_user_compliance="$7" kc_user_risk="$8" kc_user_developer="$9"

  info "Waiting for the LoadBalancer to receive an external hostname (up to 10 min)..."
  local lb_hostname=""
  local retries=60
  for ((i = 1; i <= retries; i++)); do
    lb_hostname=$(kubectl get svc -n "$NAMESPACE" "pacific-ai-ingress-nginx-controller" \
      -o jsonpath='{.status.loadBalancer.ingress[0].hostname}' 2>/dev/null || true)
    if [[ -n "$lb_hostname" ]]; then
      echo ""
      BASE_URL="http://${lb_hostname}"
      success "LoadBalancer URL: ${BASE_URL}"

      # Update the release with the real base URL
      info "Updating Helm release with the real base URL..."
      run helm upgrade pacific-ai "${chart_dir}" \
        --namespace "$NAMESPACE" \
        --reuse-values \
        --set "global.baseUrl=${BASE_URL}"

      # Restart Keycloak and the pacific-ai deployment so they pick up the new URL
      info "Restarting Keycloak and Pacific AI to apply the new URL..."
      run kubectl rollout restart statefulset pacific-ai-keycloak -n "$NAMESPACE"
      run kubectl rollout restart deployment pacific-ai -n "$NAMESPACE"
      run kubectl rollout status statefulset pacific-ai-keycloak -n "$NAMESPACE" --timeout=3m
      run kubectl rollout status deployment pacific-ai -n "$NAMESPACE" --timeout=3m

      # Re-save credentials with the real URL
      _save_credentials "$kc_admin_pw" \
        "$kc_user_admin" "$kc_user_member" "$kc_user_gov" "$kc_user_policy" \
        "$kc_user_vendor" "$kc_user_compliance" "$kc_user_risk" "$kc_user_developer"
      return
    fi
    printf "  [%d/%d] Still waiting...\r" "$i" "$retries"
    sleep 10
  done
  echo ""
  warn "Could not auto-detect the LoadBalancer hostname."
  warn "Once it's available, run:"
  warn "  helm upgrade pacific-ai ${chart_dir} --namespace ${NAMESPACE} --reuse-values --set global.baseUrl=http://<hostname>"
  warn "  kubectl rollout restart statefulset pacific-ai-keycloak -n ${NAMESPACE}"
  warn "  kubectl rollout restart deployment pacific-ai -n ${NAMESPACE}"
}

_save_credentials() {
  local kc_admin_pw="$1"
  local kc_user_admin="$2" kc_user_member="$3" kc_user_gov="$4" kc_user_policy="$5"
  local kc_user_vendor="$6" kc_user_compliance="$7" kc_user_risk="$8" kc_user_developer="$9"
  cat > "$CREDENTIALS_FILE" <<EOF
Pacific AI - Installation Credentials
Generated: $(date)

  Cluster:    ${CLUSTER_NAME}
  Namespace:  ${NAMESPACE}
  Region:     ${REGION}
  URL:        ${BASE_URL:-http://pending (see helm upgrade note above)}

Application Users:
  admin               password: ${kc_user_admin}
  member              password: ${kc_user_member}
  governance_officer  password: ${kc_user_gov}
  policy_manager      password: ${kc_user_policy}
  vendor_manager      password: ${kc_user_vendor}
  compliance_officer  password: ${kc_user_compliance}
  risk_manager        password: ${kc_user_risk}
  developer           password: ${kc_user_developer}

Infrastructure (internal use):
  Keycloak admin password: ${kc_admin_pw}

IMPORTANT: Store this file securely and delete it after noting
           your credentials. Never share or commit this file.
EOF
  chmod 600 "$CREDENTIALS_FILE"
}

# ─────────────────────────────────────────────────────────────────────────────
# Step 10: Summary
# ─────────────────────────────────────────────────────────────────────────────
print_summary() {
  step "Installation complete"

  echo ""
  echo -e "${GREEN}${BOLD}"
  cat <<'BANNER'
╔═══════════════════════════════════════════════════════════════════╗
║                  Pacific AI is Ready!                             ║
╚═══════════════════════════════════════════════════════════════════╝
BANNER
  echo -e "${NC}"

  echo -e "  URL:    ${BOLD}${BASE_URL:-see helm upgrade note above}${NC}"
  echo ""
  echo "  Login credentials:"
  echo "    Username: admin (or any user listed in ${CREDENTIALS_FILE})"
  echo "    Password: (see ${CREDENTIALS_FILE})"
  echo ""
  echo -e "${YELLOW}  Next steps:${NC}"
  echo "    1. Open the URL above in your browser"
  echo "    2. Log in as 'admin' with the password from ${CREDENTIALS_FILE}"
  echo "    3. Complete the onboarding wizard"
  echo "    4. Delete ${CREDENTIALS_FILE} after saving the credentials securely"
  echo ""
  echo -e "${YELLOW}  Useful commands:${NC}"
  echo "    kubectl get pods -n ${NAMESPACE}          # Check pod status"
  echo "    helm list -n ${NAMESPACE}                 # Check Helm release"
  echo "    helm upgrade pacific-ai ... --reuse-values # Reconfigure"
  if [[ "$TRACK_TIME" == "true" ]]; then
    local elapsed=$(( SECONDS - START_TIME ))
    local mins=$(( elapsed / 60 ))
    local secs=$(( elapsed % 60 ))
    echo -e "  ${BOLD}Total time: ${mins}m ${secs}s${NC}"
  fi
  echo ""
}

# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────
main() {
  parse_args "$@"
  [[ "$TRACK_TIME" == "true" ]] && START_TIME=$SECONDS
  print_banner
  check_prerequisites            # Step 1
  check_aws_credentials          # Step 2
  collect_region                 # Step 3
  setup_cluster                  # Step 4
  setup_namespace                # Step 5
  setup_storage                  # Step 6
  pull_helm_chart                # Step 7
  collect_configuration          # Step 8
  install_pacific_ai             # Step 9
  attach_iam_to_service_account  # Step 10
  print_summary                  # Step 11
}

main "$@"