#!/bin/bash

################################################################################
# CFN Docker Wave Execution - Container Spawning
# Purpose: Spawn Docker containers from batching plan with tier-aware memory limits
# Version: 1.0.0
# Exit Codes: 0=success, 1=execution_error, 2=validation_error
################################################################################

set -euo pipefail

# Get script directory and source helpers
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SKILL_DIR="$SCRIPT_DIR"
LIB_DIR="$SKILL_DIR/lib"

# Source helper functions
if [[ ! -f "$LIB_DIR/docker-helpers.sh" ]]; then
  echo "ERROR: docker-helpers.sh not found at $LIB_DIR/docker-helpers.sh"
  exit 2
fi
source "$LIB_DIR/docker-helpers.sh"

################################################################################
# CONFIGURATION
################################################################################

# Default values
DEFAULT_BASE_IMAGE="claude-flow-novice:latest"
DEFAULT_WORKSPACE="/workspace"
DEFAULT_NETWORK="cfn-network"
DEFAULT_SPAWN_PARALLEL=5
DEFAULT_TASK_TIMEOUT=1800

# Memory tier mapping (from cfn-error-batching-strategy)
declare -A TIER_MEMORY=(
  [1]="${CFN_TIER_1_MEMORY:-512m}"
  [2]="${CFN_TIER_2_MEMORY:-600m}"
  [3]="${CFN_TIER_3_MEMORY:-800m}"
  [4]="${CFN_TIER_4_MEMORY:-1g}"
)

# Command line arguments
WAVE_PLAN=""
WAVE_NUMBER=""
BASE_IMAGE="$DEFAULT_BASE_IMAGE"
WORKSPACE="$DEFAULT_WORKSPACE"
NETWORK="$DEFAULT_NETWORK"
OUTPUT_FILE=""
DRY_RUN=false
SPAWN_PARALLEL="$DEFAULT_SPAWN_PARALLEL"
VERBOSE=false
TASK_TIMEOUT="$DEFAULT_TASK_TIMEOUT"

declare -a EXTRA_ENV_VARS=()
declare -a EXTRA_VOLUMES=()

################################################################################
# HELP/USAGE
################################################################################

usage() {
  cat << 'EOF'
CFN Docker Wave Execution - Container Spawning

Usage: spawn-wave.sh [OPTIONS]

Required Options:
  --wave-plan FILE              Path to batching plan JSON
  --wave-number N               Wave number to spawn

Optional Options:
  --base-image IMAGE            Docker image (default: claude-flow-novice:latest)
  --workspace PATH              Workspace mount path (default: /workspace)
  --network NAME                Docker network (default: cfn-network)
  --output FILE                 Write container manifest to file
  --environment VAR=VALUE       Additional environment variables (repeatable)
  --volume SRC:DST              Additional volume mounts (repeatable)
  --parallel N                  Max concurrent spawns (default: 5)
  --task-timeout SECONDS        Task timeout (default: 1800)
  --dry-run                     Show what would be spawned
  --verbose                     Enable verbose logging
  --help                        Show this help message

Examples:
  spawn-wave.sh --wave-plan waves.json --wave-number 1
  spawn-wave.sh --wave-plan waves.json --wave-number 1 --parallel 10
  spawn-wave.sh --wave-plan waves.json --wave-number 2 --workspace /data/ws
  spawn-wave.sh --wave-plan waves.json --wave-number 1 --dry-run --verbose

EOF
}

################################################################################
# ARGUMENT PARSING
################################################################################

while [[ $# -gt 0 ]]; do
  case $1 in
    --wave-plan)
      WAVE_PLAN="$2"
      shift 2
      ;;
    --wave-number)
      WAVE_NUMBER="$2"
      shift 2
      ;;
    --base-image)
      BASE_IMAGE="$2"
      shift 2
      ;;
    --workspace)
      WORKSPACE="$2"
      shift 2
      ;;
    --network)
      NETWORK="$2"
      shift 2
      ;;
    --output)
      OUTPUT_FILE="$2"
      shift 2
      ;;
    --environment)
      # Validate environment variable before adding (CRITICAL FIX #1)
      if ! validate_environment_variable "$2"; then
        log_error "Environment variable validation failed: $2"
        usage
        exit 2
      fi
      EXTRA_ENV_VARS+=("-e" "$2")
      shift 2
      ;;
    --volume)
      EXTRA_VOLUMES+=("-v" "$2")
      shift 2
      ;;
    --parallel)
      SPAWN_PARALLEL="$2"
      shift 2
      ;;
    --task-timeout)
      TASK_TIMEOUT="$2"
      shift 2
      ;;
    --dry-run)
      DRY_RUN=true
      shift
      ;;
    --verbose)
      VERBOSE=true
      CFN_DEBUG=true
      shift
      ;;
    --help)
      usage
      exit 0
      ;;
    *)
      log_error "Unknown option: $1"
      usage
      exit 2
      ;;
  esac
done

################################################################################
# VALIDATION
################################################################################

validate_arguments() {
  if [[ -z "$WAVE_PLAN" ]]; then
    log_error "Missing required argument: --wave-plan"
    exit 2
  fi

  if [[ -z "$WAVE_NUMBER" ]]; then
    log_error "Missing required argument: --wave-number"
    exit 2
  fi

  if ! validate_json_file "$WAVE_PLAN"; then
    log_error "Invalid wave plan: $WAVE_PLAN"
    exit 2
  fi

  if ! validate_docker_access; then
    log_error "Docker not accessible"
    exit 2
  fi

  if ! validate_jq; then
    log_error "jq not available"
    exit 2
  fi

  log_debug "Arguments validated"
}

# Get task ID from plan
get_task_id() {
  jq -r '.task_id // "unknown"' "$WAVE_PLAN"
}

# Get wave data from plan
get_wave_data() {
  local wave_num="$1"

  jq --arg wave_num "$wave_num" \
    '.waves[] | select(.wave_number == ($wave_num | tonumber))' \
    "$WAVE_PLAN"
}

# Get batch data from wave (compact JSON, one per line)
get_wave_batches() {
  local wave_data="$1"

  echo "$wave_data" | jq -c '.batches[]'
}

# Get batch count for wave
get_batch_count() {
  local wave_data="$1"

  echo "$wave_data" | jq '.batch_count // (.batches | length)'
}

################################################################################
# SPAWNING LOGIC
################################################################################

# Get memory limit for tier
get_tier_memory() {
  local tier="$1"

  if [[ -n "${TIER_MEMORY[$tier]:-}" ]]; then
    echo "${TIER_MEMORY[$tier]}"
  else
    log_warn "Unknown tier: $tier, using default 512m"
    echo "512m"
  fi
}

# Create container name from batch
create_container_name() {
  local wave_num="$1"
  local batch_id="$2"

  # Generate hash-based name for collision resistance (CRITICAL FIX #2)
  local container_name
  container_name=$(generate_safe_container_name "$wave_num" "$batch_id")

  # Verify uniqueness before returning
  if container_name_exists "$container_name"; then
    log_error "Container name collision detected: $container_name"
    log_error "Original batch_id: $batch_id"
    return 1
  fi

  echo "$container_name"
}

# Spawn single container
spawn_container() {
  local wave_num="$1"
  local batch_data="$2"

  local batch_id tier memory_limit task_prompt container_name
  batch_id=$(echo "$batch_data" | jq -r '.batch_id')
  tier=$(echo "$batch_data" | jq -r '.tier')
  memory_limit=$(get_tier_memory "$tier")
  task_prompt=$(echo "$batch_data" | jq -r '.task_prompt // "Fix errors in assigned files"')

  # SANITIZE TASK PROMPT (CRITICAL FIX #3)
  task_prompt=$(sanitize_env_value "$task_prompt")

  # Create container name with collision detection
  if ! container_name=$(create_container_name "$wave_num" "$batch_id"); then
    log_error "Failed to create unique container name for batch: $batch_id"
    return 1
  fi

  log_info "Spawning container: $container_name (tier=$tier, memory=$memory_limit)"

  if [[ "$DRY_RUN" == "true" ]]; then
    log_info "[DRY-RUN] Would spawn: docker run -d \\"
    log_info "[DRY-RUN]   --name $container_name \\"
    log_info "[DRY-RUN]   --memory $memory_limit \\"
    log_info "[DRY-RUN]   --memory-reservation $memory_limit \\"
    log_info "[DRY-RUN]   -v /workspace:/workspace:rw \\"
    log_info "[DRY-RUN]   --network $NETWORK \\"
    log_info "[DRY-RUN]   -e BATCH_ID=$batch_id \\"
    log_info "[DRY-RUN]   -e WAVE_NUMBER=$wave_num \\"
    log_info "[DRY-RUN]   -e TASK_PROMPT='$task_prompt' \\"
    log_info "[DRY-RUN]   $BASE_IMAGE"

    # Still return valid container_id for dry-run
    echo "dry-run-${batch_id}-$(date +%s)"
    return 0
  fi

  # Get task ID for labels
  local task_id
  task_id=$(get_task_id)

  # Build docker run command
  local docker_opts=(
    "run"
    "-d"
    "--name" "$container_name"
    "--memory" "$memory_limit"
    "--memory-reservation" "$memory_limit"
    "-v" "$WORKSPACE:/workspace:rw"
    "--network" "$NETWORK"
    "--label" "cfn.task.id=$task_id"
    "--label" "cfn.wave.number=$wave_num"
    "--label" "cfn.batch.id=$batch_id"
    "--label" "cfn.tier=$tier"
    "--label" "cfn.memory.limit=$memory_limit"
    "-e" "BATCH_ID=$batch_id"
    "-e" "WAVE_NUMBER=$wave_num"
    "-e" "TASK_PROMPT=$task_prompt"
    "-e" "TASK_TIMEOUT=$TASK_TIMEOUT"
  )

  # Add extra environment variables
  for i in "${!EXTRA_ENV_VARS[@]}"; do
    if [[ "${EXTRA_ENV_VARS[$i]}" == "-e" ]]; then
      docker_opts+=("-e" "${EXTRA_ENV_VARS[$((i + 1))]}")
    fi
  done

  # Add extra volumes
  for i in "${!EXTRA_VOLUMES[@]}"; do
    if [[ "${EXTRA_VOLUMES[$i]}" == "-v" ]]; then
      docker_opts+=("-v" "${EXTRA_VOLUMES[$((i + 1))]}")
    fi
  done

  docker_opts+=("$BASE_IMAGE")

  # Execute docker run
  local container_id
  if container_id=$(docker "${docker_opts[@]}" 2>&1); then
    log_success "Container spawned: $container_id ($container_name)"
    echo "$container_id"
    return 0
  else
    log_error "Failed to spawn container: $container_name"
    log_error "Error: $container_id"
    return 1
  fi
}

# Spawn wave with parallelism control
spawn_wave() {
  local wave_num="$1"

  log_info "Spawning Wave $wave_num..."

  # Get wave data
  local wave_data
  if ! wave_data=$(get_wave_data "$wave_num"); then
    log_error "Wave $wave_num not found in plan"
    return 1
  fi

  if [[ -z "$wave_data" ]]; then
    log_error "Wave $wave_num has no data"
    return 1
  fi

  # Get batch count
  local batch_count
  batch_count=$(get_batch_count "$wave_data")
  log_info "Wave $wave_num has $batch_count batches"

  # Initialize manifest
  local manifest_batches=()
  local spawned_count=0
  local failed_count=0
  local active_pids=()

  # Spawn containers with parallelism control
  while IFS= read -r batch_data; do
    [[ -z "$batch_data" ]] && continue

    # Wait if we have reached parallel limit
    while (( ${#active_pids[@]} >= SPAWN_PARALLEL )); do
      # Check for completed processes
      local still_active=()
      for pid in "${active_pids[@]}"; do
        if kill -0 "$pid" 2>/dev/null; then
          still_active+=("$pid")
        fi
      done
      active_pids=("${still_active[@]}")

      if (( ${#active_pids[@]} >= SPAWN_PARALLEL )); then
        sleep 0.1
      fi
    done

    # Spawn container in background
    (
      if container_id=$(spawn_container "$wave_num" "$batch_data"); then
        echo "$container_id"
      fi
    ) &
    active_pids+=("$!")

  done < <(get_wave_batches "$wave_data")

  # Wait for all background jobs
  log_info "Waiting for all spawning jobs to complete..."
  local all_success=true
  for pid in "${active_pids[@]}"; do
    if wait "$pid" 2>/dev/null; then
      spawned_count=$((spawned_count + 1))
    else
      failed_count=$((failed_count + 1))
      all_success=false
    fi
  done

  if [[ "$all_success" == "false" ]] && [[ $spawned_count -eq 0 ]]; then
    log_error "Some containers failed to spawn: $failed_count failed, $spawned_count spawned"
    return 1
  fi

  log_success "Successfully spawned $spawned_count containers for wave $wave_num"
  return 0
}

# Collect spawned container information
collect_container_info() {
  local wave_num="$1"
  local pattern="cfn-wave${wave_num}-*"

  log_info "Collecting container information for pattern: $pattern"

  local containers_json="[]"
  local container_count=0

  # Query running and exited containers
  while IFS= read -r container_id; do
    local container_name started_at batch_id tier memory_limit status

    container_name=$(docker inspect -f '{{.Name}}' "$container_id" | sed 's|^/||')
    started_at=$(docker inspect -f '{{.State.StartedAt}}' "$container_id")
    status=$(docker inspect -f '{{.State.Status}}' "$container_id")
    batch_id=$(docker inspect -f '{{.Config.Env}}' "$container_id" | grep -oP 'BATCH_ID=\K[^,]+' || echo "unknown")
    tier=$(echo "$batch_id" | grep -oP 'batch-\d+-tier-\K\d+' || echo "1")
    memory_limit=$(docker inspect -f '{{.HostConfig.Memory}}' "$container_id" | awk '{printf "%.0fm\n", $1/1024/1024}' || echo "unknown")

    local container_obj
    container_obj=$(jq -n \
      --arg container_id "$container_id" \
      --arg container_name "$container_name" \
      --arg batch_id "$batch_id" \
      --arg tier "$tier" \
      --arg memory_limit "$memory_limit" \
      --arg status "$status" \
      --arg started_at "$started_at" \
      '{
        id: $container_id,
        container_id: $container_id,
        container_name: $container_name,
        batch_id: $batch_id,
        tier: ($tier | tonumber),
        memory_limit: $memory_limit,
        status: $status,
        started_at: $started_at
      }')

    containers_json=$(echo "$containers_json" | jq --argjson obj "$container_obj" '. += [$obj]')
    container_count=$((container_count + 1))

  done < <(docker ps -a --filter "name=$pattern" --format "{{.ID}}")

  log_info "Found $container_count containers"

  echo "$containers_json"
}

################################################################################
# MAIN
################################################################################

main() {
  # Validate arguments
  validate_arguments

  # Create network if needed
  if ! create_network_if_missing "$NETWORK"; then
    log_warn "Failed to create network, attempting to proceed anyway"
  fi

  # Validate base image exists or can be pulled
  log_info "Checking Docker image: $BASE_IMAGE"
  if ! docker inspect "$BASE_IMAGE" > /dev/null 2>&1; then
    log_warn "Image not found locally, attempting to pull: $BASE_IMAGE"
    if ! docker pull "$BASE_IMAGE"; then
      log_error "Failed to pull image: $BASE_IMAGE"
      exit 2
    fi
  fi

  # Spawn wave
  if ! spawn_wave "$WAVE_NUMBER"; then
    log_error "Wave spawning failed"
    exit 1
  fi

  # Collect container information
  local container_info
  container_info=$(collect_container_info "$WAVE_NUMBER")

  # Build output manifest
  local manifest
  manifest=$(jq -n \
    --arg wave_number "$WAVE_NUMBER" \
    --arg spawned_at "$(date -u +'%Y-%m-%dT%H:%M:%SZ')" \
    --argjson containers "$container_info" \
    --arg total_spawned "$(echo "$container_info" | jq 'length')" \
    '{
      wave_number: ($wave_number | tonumber),
      spawned_at: $spawned_at,
      containers: $containers,
      total_spawned: ($total_spawned | tonumber),
      total_memory: "calculated"
    }')

  # Save to output file if requested
  if [[ -n "$OUTPUT_FILE" ]]; then
    echo "$manifest" > "$OUTPUT_FILE"
    log_success "Container manifest saved to: $OUTPUT_FILE"
  fi

  # Output manifest to stdout
  echo "$manifest"
  log_success "Wave $WAVE_NUMBER spawning complete"
  exit 0
}

# Run main function
main