#!/bin/sh
set -eu

# ============================================================================
# Remote Coverage Agent (JaCoCo) - POSIX sh (busybox/dash compatible)
#
# AUTO start:
#   - finds the real Java PID for <processKey> (e.g. /app/app.jar)
#   - auto-detects destfile from -javaagent:...=destfile=...
#   - truncates destfile (note: output=file truncation does NOT reset in-memory,
#     but since we KILL on stop, you still get a clean exec for that run)
#
# AUTO stop:
#   - finds PID + destfile BEFORE kill
#   - derives includeCSV from cmdline if includeCSV=auto
#   - SIGTERM then SIGKILL (flush output=file)
#   - generates report using jacococli report (with Spring Boot jar handling)
# ============================================================================

ACTION="${1:-}"
MODE="${2:-}"

TOOLS_DIR="${TOOLS_DIR:-$HOME/.coverage_tools}"
JACOCOCLI="${JACOCOCLI:-$TOOLS_DIR/jacococli.jar}"

AUTO_STOP_KILL="${AUTO_STOP_KILL:-1}"     # keep as 1 for your requirement
POST_STOP_SLEEP="${POST_STOP_SLEEP:-10}"

log() { echo "[$(date +'%F %T')] $*"; }
err() { echo "[$(date +'%F %T')] ERROR: $*" >&2; }
has_cmd() { command -v "$1" >/dev/null 2>&1; }

validate_mode() {
  case "$MODE" in
    manual|auto|explicit) : ;;
    *) err "Invalid mode '$MODE' (allowed: manual|auto|explicit)"; exit 1 ;;
  esac
}

# Read /proc/<pid>/cmdline into newline-separated args (robust on busybox)
proc_args_nl() {
  pid="$1"
  [ -r "/proc/$pid/cmdline" ] || { echo ""; return 0; }
  tr '\000' '\n' < "/proc/$pid/cmdline" 2>/dev/null || true
}

# Only consider real java processes (avoid matching sh -lc wrappers)
is_java_pid() {
  pid="$1"
  first="$(proc_args_nl "$pid" | sed -n '1p')"
  [ -n "$first" ] || return 1
  case "$first" in
    java|*"/java"|*"/java.exe"|*"/jre/bin/java"|*"/bin/java") return 0 ;;
  esac
  return 1
}

# Avoid matching our own script
is_our_wrapper() {
  pid="$1"
  proc_args_nl "$pid" | grep -q 'remote_covk8s_agent\.sh' 2>/dev/null && return 0
  return 1
}

# Resolve PID from:
#  - numeric PID
#  - key="auto": a java process containing destfile= in its -javaagent opts
#  - jar path: match exact token after "-jar"
#  - fallback: substring match within java args (still java-only)
resolve_target_pid() {
  key="$1"
  [ -n "$key" ] || { err "process key is empty"; return 1; }

  # numeric PID?
  case "$key" in
    *[!0-9]* ) : ;;
    * )
      kill -0 "$key" 2>/dev/null && { echo "$key"; return 0; }
      err "PID $key is not running"
      return 1
      ;;
  esac

  hits=""
  for d in /proc/[0-9]*; do
    pid="${d#/proc/}"
    is_java_pid "$pid" || continue
    is_our_wrapper "$pid" && continue

    args="$(proc_args_nl "$pid")"
    [ -n "$args" ] || continue

    if [ "$key" = "auto" ]; then
      # java with jacoco destfile
      echo "$args" | grep -q '^-javaagent:.*jacoco.*destfile=' 2>/dev/null && hits="$hits $pid"
      continue
    fi

    case "$key" in
      *.jar|*/*.jar|/app/app.jar)
        # exact: "-jar" then next token equals key
        jar_ok="$(echo "$args" | awk -v want="$key" '
          $0=="-jar" {getline n; if(n==want){print "yes"}; exit}
        ')"
        [ "$jar_ok" = "yes" ] && hits="$hits $pid"
        ;;
      *)
        echo "$args" | grep -Fq "$key" 2>/dev/null && hits="$hits $pid"
        ;;
    esac
  done

  count="$(echo "$hits" | awk '{print NF}')"
  if [ "$count" -eq 1 ]; then
    echo "$hits" | awk '{print $1}'
    return 0
  fi
  if [ "$count" -gt 1 ]; then
    err "Multiple java PIDs match '$key':$hits"
    return 1
  fi

  err "No java PID matches '$key'"
  return 1
}

# Extract jacoco option string from arg:
# -javaagent:/jacoco/jacoco-agent.jar=destfile=/jacoco/jacoco.exec,output=file
get_jacoco_opts() {
  pid="$1"
  args="$(proc_args_nl "$pid")"
  [ -n "$args" ] || { echo ""; return 0; }
  echo "$args" | awk '
    index($0,"-javaagent:")==1 && index($0,"jacoco")>0 {
      pos=index($0,"="); if(pos>0){print substr($0,pos+1)}; exit
    }
  '
}

# From "a=b,c=d" get value for name
get_opt_val() {
  opts="$1"
  name="$2"
  echo "$opts" | awk -v n="$name" '
    BEGIN{FS=","}
    {
      for(i=1;i<=NF;i++){
        split($i,a,"=")
        if(a[1]==n){print a[2]; exit}
      }
    }
  '
}

detect_jacoco_exec_from_pid() {
  pid="$1"
  opts="$(get_jacoco_opts "$pid")"
  [ -n "$opts" ] || { echo ""; return 0; }
  get_opt_val "$opts" "destfile"
}

empty_exec_file() {
  path="$1"
  [ -n "$path" ] || { err "empty_exec_file: path empty"; return 1; }
  dir="$(dirname "$path")"
  mkdir -p "$dir" 2>/dev/null || true
  if : >"$path" 2>/dev/null; then
    log "✅ Emptied JaCoCo exec file: $path"
    return 0
  fi
  err "WARN: could not truncate $path (permission denied). Continuing."
  return 0
}

derive_include_from_cmdline_pid() {
  pid="$1"
  args="$(proc_args_nl "$pid")"
  [ -n "$args" ] || { echo ""; return 0; }

  jar_path="$(echo "$args" | awk '$0=="-jar"{getline n; print n; exit}')"
  if [ -n "${jar_path:-}" ]; then
    echo "$jar_path"
    return 0
  fi

  cp_path="$(echo "$args" | awk '$0=="-cp" || $0=="-classpath"{getline n; print n; exit}')"
  if [ -n "${cp_path:-}" ]; then
    echo "$cp_path" | tr ':' ','
    return 0
  fi

  echo ""
}

# --- Spring Boot fat jar handling (avoid BOOT-INF/lib duplicate classes) ---
is_spring_boot_jar() {
  j="$1"
  [ -f "$j" ] || return 1
  if has_cmd jar; then
    jar tf "$j" 2>/dev/null | awk '
      $0 ~ /^BOOT-INF\/classes\// { found=1; exit }
      END { exit(found?0:1) }
    '
    return $?
  fi
  if has_cmd unzip; then
    unzip -l "$j" 2>/dev/null | awk '
      $0 ~ /BOOT-INF\/classes\// { found=1; exit }
      END { exit(found?0:1) }
    '
    return $?
  fi
  return 1
}

extract_boot_classes() {
  j="$1"
  dest="$2"
  rm -rf "$dest" 2>/dev/null || true
  mkdir -p "$dest" 2>/dev/null || true

  if has_cmd jar; then
    ( cd "$dest" && jar xf "$j" BOOT-INF/classes ) >/dev/null 2>&1 || return 1
    [ -d "$dest/BOOT-INF/classes" ] || return 1
    echo "$dest/BOOT-INF/classes"
    return 0
  fi

  if has_cmd unzip; then
    unzip -oq "$j" 'BOOT-INF/classes/*' -d "$dest" >/dev/null 2>&1 || return 1
    [ -d "$dest/BOOT-INF/classes" ] || return 1
    echo "$dest/BOOT-INF/classes"
    return 0
  fi

  return 1
}

normalize_include_csv() {
  include_csv="$1"
  uuid="$2"
  [ -n "${include_csv:-}" ] || { echo ""; return 0; }

  # clean spaces around commas
  include_csv="$(echo "$include_csv" | sed 's/[[:space:]]*,[[:space:]]*/,/g')"

  out=""
  oldifs="$IFS"; IFS=','

  for p in $include_csv; do
    p="$(printf "%s" "$p" | awk '{$1=$1;print}')"
    [ -n "$p" ] || continue

    case "$p" in
      *.jar|*/*.jar)
        if is_spring_boot_jar "$p"; then
          tmpdir="$TOOLS_DIR/boot_classes_$uuid"
          cls="$(extract_boot_classes "$p" "$tmpdir" 2>/dev/null || true)"
          if [ -n "${cls:-}" ] && [ -d "$cls" ]; then
            log "✅ Spring Boot jar detected. Using classfiles: $cls"
            [ -z "$out" ] && out="$cls" || out="$out,$cls"
          else
            err "Spring Boot jar detected but cannot extract BOOT-INF/classes."
            err "Install 'jar' (JDK) or 'unzip' in container OR generate report on controller."
            exit 1
          fi
        else
          log "INFO: Skipping non-Spring JAR from classfiles: $p"  
      	#	[ -z "$out" ] && out="$p" || out="$out,$p"
        fi
        ;;
      *)
        [ -z "$out" ] && out="$p" || out="$out,$p"
        ;;
    esac
  done

  IFS="$oldifs"
  echo "$out"
}

gen_report() {
  exec_path="$1"
  include_csv="$2"
  uuid="$3"
  source_csv="${4:-}"

  [ -f "$JACOCOCLI" ] || { err "jacococli.jar not found at $JACOCOCLI"; exit 1; }
  [ -f "$exec_path" ] || { err "JaCoCo exec file not found: $exec_path"; exit 1; }
  has_cmd java || { err "java not found in container"; exit 1; }

  report_root="$TOOLS_DIR/report-$uuid"
  html_dir="$report_root/html"
  xml_file="$report_root/report.xml"
  csv_file="$report_root/report.csv"
  exec_copy="$TOOLS_DIR/report_${uuid}.exec"

  rm -rf "$report_root" 2>/dev/null || true
  mkdir -p "$html_dir" 2>/dev/null || true

  cp "$exec_path" "$exec_copy" 2>/dev/null || true
  [ -f "$exec_copy" ] && log "Copied exec file to $exec_copy" || log "WARN: could not copy exec to $exec_copy"

  include_csv="$(normalize_include_csv "$include_csv" "$uuid")"
  [ -n "$include_csv" ] || { err "includeCSV empty"; exit 1; }

  set -- report "$exec_path"

  oldifs="$IFS"; IFS=','

  for p in $include_csv; do
    [ -n "$p" ] || continue
    if [ -f "$p" ]; then
      case "$p" in
        *.class) set -- "$@" --classfiles "$(dirname "$p")" ;;
        *)       set -- "$@" --classfiles "$p" ;;
      esac
    elif [ -d "$p" ]; then
      set -- "$@" --classfiles "$p"
    else
      log "WARN: include path not found: $p"
    fi
  done

  IFS="$oldifs"

  src_norm="$(echo "${source_csv:-}" | sed 's/[[:space:]]*,[[:space:]]*/,/g')"
  if [ -n "$src_norm" ]; then
    oldifs="$IFS"; IFS=','

    for s in $src_norm; do
      [ -n "$s" ] || continue
      [ -d "$s" ] && set -- "$@" --sourcefiles "$s" || log "WARN: source path not found: $s"
    done

    IFS="$oldifs"
  fi

#  set -- "$@" --html "$html_dir" --xml "$xml_file" --csv "$csv_file"

 # log "Generating JaCoCo report (uuid=$uuid)..."
  #java -jar "$JACOCOCLI" "$@"
  set -- "$@" \
   --exclude "**/BOOT-INF/lib/**" \
   --exclude "**/lombok/**" \
   --exclude "**/org/projectlombok/**" \
   --html "$html_dir" \
   --xml "$xml_file" \
   --csv "$csv_file"

  log "Generating JaCoCo report (uuid=$uuid)..."
  java -jar "$JACOCOCLI" "$@"

  ( cd "$html_dir" && tar -czf "report_${uuid}.tar.gz" . ) >/dev/null 2>&1 || true
  log "✅ Report ready: $html_dir/report_${uuid}.tar.gz"
}

graceful_stop_pid() {
  pid="$1"
  [ -n "${pid:-}" ] || return 0
  if ! kill -0 "$pid" 2>/dev/null; then
    log "No such PID: $pid (already stopped?)"
    return 0
  fi
  log "Stopping PID $pid (SIGTERM)..."
  kill -15 "$pid" 2>/dev/null || true
  sleep 8
  if kill -0 "$pid" 2>/dev/null; then
    log "PID $pid still alive, forcing SIGKILL..."
    kill -9 "$pid" 2>/dev/null || true
  fi
}

validate_mode

case "$ACTION" in
  start)
    # start <mode> <processKey> <jacocoExecPathOrEmpty> <startCmd> <stopCmdPlaceholder>
    PROCESS_KEY="${3:-auto}"
    JACOCO_EXEC_ARG="${4:-}"
    START_CMD="${5:-}"

    pid="$(resolve_target_pid "$PROCESS_KEY" 2>/dev/null || true)"
    [ -n "$pid" ] || { err "Auto start: process '$PROCESS_KEY' not running (PID not found)"; exit 1; }

    if [ -n "$JACOCO_EXEC_ARG" ]; then
      JACOCO_EXEC="$JACOCO_EXEC_ARG"
    else
      JACOCO_EXEC="$(detect_jacoco_exec_from_pid "$pid")"
    fi

    [ -n "${JACOCO_EXEC:-}" ] || { err "Auto start: could not detect destfile (no jacoco -javaagent?)"; exit 1; }

    empty_exec_file "$JACOCO_EXEC" || true

    # Optional explicit start command (usually not needed in your flow)
    if [ "$MODE" = "explicit" ] && [ -n "${START_CMD:-}" ]; then
      log "Running explicit start command..."
      sh -c "$START_CMD"
      log "Explicit start command finished."
    fi

    log "✅ Start complete (pid=$pid destfile=$JACOCO_EXEC)."
    ;;

  stop)
    # stop <mode> <processKey> <jacocoExecPathOrEmpty> <includeCSV|auto> <uuid> <stopCmd> <sourceCSV>
    PROCESS_KEY="${3:-auto}"
    JACOCO_EXEC_ARG="${4:-}"
    INCLUDE_CSV="${5:-auto}"
    UUID="${6:-}"
    STOP_CMD="${7:-}"
    SOURCE_CSV="${8:-}"

    [ -n "$UUID" ] || { err "uuid is required for stop"; exit 1; }

    pid="$(resolve_target_pid "$PROCESS_KEY" 2>/dev/null || true)"
    [ -n "$pid" ] || { err "Auto stop: process '$PROCESS_KEY' not running (PID not found)"; exit 1; }

    # detect destfile BEFORE killing (critical)
    if [ -n "$JACOCO_EXEC_ARG" ]; then
      JACOCO_EXEC="$JACOCO_EXEC_ARG"
    else
      JACOCO_EXEC="$(detect_jacoco_exec_from_pid "$pid")"
    fi
    [ -n "${JACOCO_EXEC:-}" ] || { err "Auto stop: could not detect destfile (no jacoco -javaagent?)"; exit 1; }

    # derive include BEFORE killing (critical)
    if [ -z "${INCLUDE_CSV:-}" ] || [ "$INCLUDE_CSV" = "auto" ]; then
      INCLUDE_CSV="$(derive_include_from_cmdline_pid "$pid")"
      [ -n "${INCLUDE_CSV:-}" ] || {
        err "includeCSV=auto but could not derive (-jar/-cp not found). Provide includeCSV manually."
        exit 1
      }
      log "✅ Derived includeCSV: $INCLUDE_CSV"
    fi

    # you requested: always kill on stop
    if [ "$AUTO_STOP_KILL" = "1" ]; then
      graceful_stop_pid "$pid"
    else
      log "AUTO_STOP_KILL=0: not killing process (but output=file may not flush)."
    fi

    sleep "$POST_STOP_SLEEP"
    log "Using JaCoCo exec: $JACOCO_EXEC"
    gen_report "$JACOCO_EXEC" "$INCLUDE_CSV" "$UUID" "$SOURCE_CSV"
    log "✅ Stop complete."
    ;;

  *)
    err "Invalid action (use start|stop)"
    exit 1
    ;;
esac