#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'EOF'
Usage:
  scripts/ci/enrich_exploitability.sh --mode single|batch --input <path> --output <path> [--cvss-vectors <path>] [--analyzer <path>]

Options:
  --mode           Processing mode: single advisory object or batch advisory array
  --input          Input JSON path
  --output         Output JSON path
  --cvss-vectors   Optional JSON object mapping advisory id -> CVSS vector
  --analyzer       Optional analyzer path (default: utils/analyze_exploitability.py)
  --help           Show this help
EOF
}

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "$REPO_ROOT"

MODE=""
INPUT_PATH=""
OUTPUT_PATH=""
CVSS_VECTORS_PATH=""
ANALYZER_PATH="utils/analyze_exploitability.py"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --mode)
      MODE="${2:-}"
      shift 2
      ;;
    --input)
      INPUT_PATH="${2:-}"
      shift 2
      ;;
    --output)
      OUTPUT_PATH="${2:-}"
      shift 2
      ;;
    --cvss-vectors)
      CVSS_VECTORS_PATH="${2:-}"
      shift 2
      ;;
    --analyzer)
      ANALYZER_PATH="${2:-}"
      shift 2
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      echo "ERROR: Unknown argument: $1" >&2
      usage >&2
      exit 1
      ;;
  esac
done

if [[ "$MODE" != "single" && "$MODE" != "batch" ]]; then
  echo "ERROR: --mode must be one of: single, batch" >&2
  exit 1
fi

if [[ -z "$INPUT_PATH" || -z "$OUTPUT_PATH" ]]; then
  echo "ERROR: --input and --output are required" >&2
  exit 1
fi

if [[ ! -f "$INPUT_PATH" ]]; then
  echo "ERROR: input file not found: $INPUT_PATH" >&2
  exit 1
fi

if [[ ! -f "$ANALYZER_PATH" ]]; then
  echo "ERROR: analyzer file not found: $ANALYZER_PATH" >&2
  exit 1
fi

if [[ -n "$CVSS_VECTORS_PATH" && ! -f "$CVSS_VECTORS_PATH" ]]; then
  echo "ERROR: --cvss-vectors file not found: $CVSS_VECTORS_PATH" >&2
  exit 1
fi

if ! command -v jq >/dev/null 2>&1; then
  echo "ERROR: jq is required" >&2
  exit 1
fi

if command -v python >/dev/null 2>&1; then
  PYTHON_BIN="python"
elif command -v python3 >/dev/null 2>&1; then
  PYTHON_BIN="python3"
else
  echo "ERROR: python or python3 is required" >&2
  exit 1
fi

tmpdir="$(mktemp -d)"
trap 'rm -rf "$tmpdir"' EXIT

resolve_cvss_vector() {
  local advisory_json="$1"
  local advisory_id
  advisory_id="$(echo "$advisory_json" | jq -r '.id // ""')"

  if [[ -n "$CVSS_VECTORS_PATH" ]]; then
    jq -r --arg id "$advisory_id" '.[$id] // ""' "$CVSS_VECTORS_PATH"
  else
    echo "$advisory_json" | jq -r '.cvss_vector // ""'
  fi
}

severity_to_cvss() {
  case "$1" in
    critical) echo "9.5" ;;
    high) echo "7.5" ;;
    medium) echo "5.5" ;;
    low) echo "3.0" ;;
    *) echo "5.0" ;;
  esac
}

build_analysis_input() {
  local advisory_json="$1"
  local mode="$2"
  local cve_id cvss_score cvss_vector vuln_type description references severity

  cve_id="$(echo "$advisory_json" | jq -r '.id // ""')"
  vuln_type="$(echo "$advisory_json" | jq -r '.type // ""')"
  description="$(echo "$advisory_json" | jq -r '.description // ""')"
  references="$(echo "$advisory_json" | jq -c '.references // []')"
  cvss_vector="$(resolve_cvss_vector "$advisory_json")"

  if [[ "$mode" == "single" ]]; then
    severity="$(echo "$advisory_json" | jq -r '.severity // "medium"')"
    cvss_score="$(severity_to_cvss "$severity")"
  else
    cvss_score="$(echo "$advisory_json" | jq -r '.cvss_score // 0')"
  fi

  jq -n \
    --arg cve_id "$cve_id" \
    --argjson cvss_score "$cvss_score" \
    --arg cvss_vector "$cvss_vector" \
    --arg type "$vuln_type" \
    --arg description "$description" \
    --argjson references "$references" \
    '{
      cve_id: $cve_id,
      cvss_score: $cvss_score,
      cvss_vector: $cvss_vector,
      type: $type,
      description: $description,
      references: $references
    }'
}

run_analysis() {
  local advisory_json="$1"
  local mode="$2"
  local output_file="$3"
  local advisory_id analysis_input analysis

  advisory_id="$(echo "$advisory_json" | jq -r '.id // "unknown"')"
  analysis_input="$(build_analysis_input "$advisory_json" "$mode")"

  if analysis="$(echo "$analysis_input" | "$PYTHON_BIN" "$ANALYZER_PATH" --json --check-exploits 2>/dev/null)"; then
    echo "$analysis" > "$output_file"
    return 0
  fi

  echo "::warning::Failed to analyze exploitability for $advisory_id, continuing without enrichment"
  return 1
}

enrich_single() {
  if ! jq -e 'type == "object"' "$INPUT_PATH" >/dev/null; then
    echo "ERROR: single mode expects JSON object at $INPUT_PATH" >&2
    exit 1
  fi

  local advisory analysis_file output_tmp
  advisory="$(cat "$INPUT_PATH")"
  analysis_file="$tmpdir/analysis_single.json"
  output_tmp="$tmpdir/output_single.json"

  if run_analysis "$advisory" "single" "$analysis_file"; then
    jq --slurpfile analysis "$analysis_file" '
      . + {
        exploitability_score: $analysis[0].exploitability_score,
        exploitability_rationale: $analysis[0].exploitability_rationale,
        attack_vector_analysis: $analysis[0].attack_vector_analysis,
        exploit_detection: $analysis[0].exploit_detection
      }
    ' "$INPUT_PATH" > "$output_tmp"
  else
    cp "$INPUT_PATH" "$output_tmp"
  fi

  mv "$output_tmp" "$OUTPUT_PATH"
  echo "Exploitability enrichment complete (single): $OUTPUT_PATH"
}

enrich_batch() {
  if ! jq -e 'type == "array"' "$INPUT_PATH" >/dev/null; then
    echo "ERROR: batch mode expects JSON array at $INPUT_PATH" >&2
    exit 1
  fi

  local analyzed_count failed_count index advisory analysis_file output_tmp analyses_json
  analyzed_count=0
  failed_count=0
  index=0
  analyses_json="$tmpdir/analyses.json"
  output_tmp="$tmpdir/output_batch.json"

  while IFS= read -r advisory; do
    analysis_file="$tmpdir/analysis_${index}.json"
    if run_analysis "$advisory" "batch" "$analysis_file"; then
      analyzed_count=$((analyzed_count + 1))
    else
      failed_count=$((failed_count + 1))
      rm -f "$analysis_file"
    fi
    index=$((index + 1))
  done < <(jq -c '.[]' "$INPUT_PATH")

  if ls "$tmpdir"/analysis_*.json >/dev/null 2>&1; then
    jq -s '.' "$tmpdir"/analysis_*.json > "$analyses_json"
  else
    echo '[]' > "$analyses_json"
  fi

  jq --slurpfile analyses "$analyses_json" '
    map(
      . as $advisory |
      ($analyses[0] | map(select(.cve_id == $advisory.id)) | first) as $analysis |
      if $analysis then
        $advisory + {
          exploitability_score: $analysis.exploitability_score,
          exploitability_rationale: $analysis.exploitability_rationale,
          attack_vector_analysis: $analysis.attack_vector_analysis,
          exploit_detection: $analysis.exploit_detection
        }
      else
        $advisory
      end
    )
  ' "$INPUT_PATH" > "$output_tmp"

  mv "$output_tmp" "$OUTPUT_PATH"
  echo "Exploitability enrichment complete (batch): $OUTPUT_PATH"
  echo "Analyzed: $analyzed_count, failed: $failed_count"
}

if [[ "$MODE" == "single" ]]; then
  enrich_single
else
  enrich_batch
fi