#!/usr/bin/env bash

#########################################################################################
#### opa-policy-tester.sh — Run OPA/Rego policy bundles against infrastructure       ####
#### configs, report violations with severity levels and CI-friendly exit codes       ####
#### Requires: bash 4+, opa binary (auto-install available)                          ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./opa-policy-tester.sh --eval --policy ./policies --input config.json         ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
OPA_PATH="${OPA_PATH:-opa}"
POLICY_DIR="${OPA_POLICY_DIR:-}"
INPUT_PATH="${OPA_INPUT:-}"
DATA_FILE=""
QUERY_DENY="${OPA_QUERY:-data.main.deny}"
QUERY_WARN="data.main.warn"
FAIL_ON="${OPA_FAIL_ON:-deny}"
OUTPUT_FORMAT="${OPA_FORMAT:-text}"
OUTPUT_FILE=""
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
RUN_MODE=""
START_TIME=""
TOTAL_INPUTS=0
TOTAL_POLICIES=0
TOTAL_PASS=0
TOTAL_WARN=0
TOTAL_FAIL=0
TMPDIR_WORK=""
RESULTS_JSON="[]"

# ── Colors ────────────────────────────────────────────────────────────
RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""

setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }
die()     { err "$*"; exit 1; }

# ── Helpers ───────────────────────────────────────────────────────────
section_header() {
    echo ""
    echo -e "  ${BOLD}${CYAN}── $1 ──${RESET}"
    echo ""
}

field() {
    printf "  ${BOLD}%-22s${RESET} %s\n" "$1" "$2"
}

field_color() {
    printf "  ${BOLD}%-22s${RESET} %b\n" "$1" "$2"
}

elapsed() {
    local end_time
    end_time=$(date +%s)
    echo "$(( end_time - START_TIME ))s"
}

cleanup() {
    if [[ -n "$TMPDIR_WORK" && -d "$TMPDIR_WORK" ]]; then
        rm -rf "$TMPDIR_WORK"
    fi
}
trap cleanup EXIT

make_tmpdir() {
    if [[ -z "$TMPDIR_WORK" ]]; then
        TMPDIR_WORK=$(mktemp -d "${TMPDIR:-/tmp}/opa-tester.XXXXXX")
    fi
}

# ── Dependency checks ────────────────────────────────────────────────
require_opa() {
    if ! command -v "$OPA_PATH" &>/dev/null; then
        err "OPA binary not found: ${OPA_PATH}"
        err "Install with:  ${SCRIPT_NAME} --install"
        err "Or set OPA_PATH to the opa binary location"
        exit 1
    fi
    verbose "OPA found: $(command -v "$OPA_PATH") ($("$OPA_PATH" version 2>/dev/null | head -1 || echo 'unknown'))"
}

require_jq() {
    if ! command -v jq &>/dev/null; then
        die "jq is required but not installed"
    fi
}

yaml_to_json() {
    local yaml_file="$1" json_file="$2"
    if command -v yq &>/dev/null; then
        yq -o=json '.' "$yaml_file" > "$json_file"
    elif command -v python3 &>/dev/null; then
        python3 -c "
import sys, json, yaml
with open('$yaml_file') as f:
    data = yaml.safe_load(f)
json.dump(data, sys.stdout)
" > "$json_file"
    else
        die "Cannot convert YAML to JSON — install yq or python3 with PyYAML"
    fi
}

# ── Install OPA ──────────────────────────────────────────────────────
install_opa() {
    local os arch url dest
    os="$(uname -s | tr '[:upper:]' '[:lower:]')"
    arch="$(uname -m)"
    case "$arch" in
        x86_64)  arch="amd64" ;;
        aarch64|arm64) arch="arm64" ;;
        *) die "Unsupported architecture: $arch" ;;
    esac

    dest="/usr/local/bin/opa"
    if [[ ! -w "$(dirname "$dest")" ]]; then
        dest="${HOME}/.local/bin/opa"
        mkdir -p "$(dirname "$dest")"
    fi

    url="https://github.com/open-policy-agent/opa/releases/latest/download/opa_${os}_${arch}"
    log "Downloading OPA from ${url}"
    if command -v curl &>/dev/null; then
        curl -fsSL -o "$dest" "$url"
    elif command -v wget &>/dev/null; then
        wget -qO "$dest" "$url"
    else
        die "Neither curl nor wget available for download"
    fi
    chmod +x "$dest"
    log "OPA installed to ${dest}"
    "$dest" version
}

# ── Collect input files ──────────────────────────────────────────────
collect_inputs() {
    local path="$1"
    local -n arr=$2
    if [[ -f "$path" ]]; then
        arr+=("$path")
    elif [[ -d "$path" ]]; then
        while IFS= read -r -d '' f; do
            arr+=("$f")
        done < <(find "$path" -type f \( -name '*.json' -o -name '*.yaml' -o -name '*.yml' \) -print0 | sort -z)
        [[ ${#arr[@]} -eq 0 ]] && die "No JSON/YAML files found in ${path}"
    else
        die "Input path does not exist: ${path}"
    fi
}

prepare_input() {
    local file="$1"
    case "$file" in
        *.yaml|*.yml)
            make_tmpdir
            local tmp_json
            tmp_json="${TMPDIR_WORK}/$(basename "${file}").json"
            verbose "Converting YAML to JSON: ${file}"
            yaml_to_json "$file" "$tmp_json"
            echo "$tmp_json"
            ;;
        *.json)
            echo "$file"
            ;;
        *)
            die "Unsupported file format: ${file}"
            ;;
    esac
}

# ── Count .rego files ────────────────────────────────────────────────
count_policies() {
    local path="$1"
    if [[ -f "$path" ]]; then
        echo 1
    elif [[ -d "$path" ]]; then
        find "$path" -name '*.rego' -not -name '*_test.rego' | wc -l | tr -d ' '
    else
        echo 0
    fi
}

# ── Evaluate policies ───────────────────────────────────────────────
eval_single() {
    local input_file="$1" json_input policy_args raw_deny raw_warn
    json_input=$(prepare_input "$input_file")

    policy_args=()
    if [[ -d "$POLICY_DIR" ]]; then
        policy_args+=(--bundle "$POLICY_DIR")
    else
        policy_args+=(--data "$POLICY_DIR")
    fi
    [[ -n "$DATA_FILE" ]] && policy_args+=(--data "$DATA_FILE")

    verbose "Evaluating: ${input_file}"
    verbose "  policy: ${POLICY_DIR}"
    verbose "  query deny: ${QUERY_DENY}"
    verbose "  query warn: ${QUERY_WARN}"

    raw_deny=$("$OPA_PATH" eval "${policy_args[@]}" --input "$json_input" \
        --format json "$QUERY_DENY" 2>/dev/null || echo '{"result":[]}')
    raw_warn=$("$OPA_PATH" eval "${policy_args[@]}" --input "$json_input" \
        --format json "$QUERY_WARN" 2>/dev/null || echo '{"result":[]}')

    local deny_msgs warn_msgs
    deny_msgs=$(echo "$raw_deny" | jq -r '
        [.result[]?.expressions[]?.value // [] | if type == "array" then .[] else . end] |
        if length > 0 then .[] else empty end' 2>/dev/null || true)
    warn_msgs=$(echo "$raw_warn" | jq -r '
        [.result[]?.expressions[]?.value // [] | if type == "array" then .[] else . end] |
        if length > 0 then .[] else empty end' 2>/dev/null || true)

    local file_denies=0 file_warns=0
    local input_base
    input_base=$(basename "$input_file")

    if [[ -n "$deny_msgs" ]]; then
        while IFS= read -r msg; do
            [[ -z "$msg" ]] && continue
            ((file_denies++)) || true
            ((TOTAL_FAIL++)) || true
            RESULTS_JSON=$(echo "$RESULTS_JSON" | jq --arg f "$input_base" \
                --arg m "$msg" --arg s "deny" \
                '. + [{"file": $f, "severity": $s, "message": $m}]')
        done <<< "$deny_msgs"
    fi

    if [[ -n "$warn_msgs" ]]; then
        while IFS= read -r msg; do
            [[ -z "$msg" ]] && continue
            ((file_warns++)) || true
            ((TOTAL_WARN++)) || true
            RESULTS_JSON=$(echo "$RESULTS_JSON" | jq --arg f "$input_base" \
                --arg m "$msg" --arg s "warn" \
                '. + [{"file": $f, "severity": $s, "message": $m}]')
        done <<< "$warn_msgs"
    fi

    if [[ $file_denies -eq 0 && $file_warns -eq 0 ]]; then
        ((TOTAL_PASS++)) || true
    fi
}

run_eval() {
    [[ -z "$POLICY_DIR" ]] && die "No policy path specified (--policy)"
    [[ -z "$INPUT_PATH" ]] && die "No input path specified (--input)"
    [[ ! -e "$POLICY_DIR" ]] && die "Policy path does not exist: ${POLICY_DIR}"
    require_opa
    require_jq

    START_TIME=$(date +%s)
    local inputs=()
    collect_inputs "$INPUT_PATH" inputs
    TOTAL_INPUTS=${#inputs[@]}
    TOTAL_POLICIES=$(count_policies "$POLICY_DIR")

    section_header "OPA Policy Evaluation"
    field "Policy path:" "$POLICY_DIR"
    field "Input path:" "$INPUT_PATH"
    field "Input files:" "$TOTAL_INPUTS"
    field "Policies:" "$TOTAL_POLICIES"
    field "Fail threshold:" "$FAIL_ON"
    echo ""

    for input_file in "${inputs[@]}"; do
        eval_single "$input_file"
    done

    render_results
    write_output
    compute_exit
}

# ── Render results ───────────────────────────────────────────────────
render_results() {
    local violation_count
    violation_count=$(echo "$RESULTS_JSON" | jq 'length')

    if [[ "$violation_count" -gt 0 ]]; then
        section_header "Violations"
        printf "  ${BOLD}%-28s %-8s %s${RESET}\n" "FILE" "LEVEL" "MESSAGE"
        printf "  %-28s %-8s %s\n" "----------------------------" "--------" "$(printf '%0.s-' {1..40})"

        echo "$RESULTS_JSON" | jq -r '.[] | [.file, .severity, .message] | @tsv' | \
        while IFS=$'\t' read -r vf vs vm; do
            local color="$RESET"
            case "$vs" in
                deny) color="$RED" ;;
                warn) color="$YELLOW" ;;
            esac
            printf "  %-28s ${color}%-8s${RESET} %s\n" "$vf" "${vs^^}" "$vm"
        done
    fi

    section_header "Summary"
    field "Total inputs:" "$TOTAL_INPUTS"
    field "Policies evaluated:" "$TOTAL_POLICIES"
    field_color "Passed:" "${GREEN}${TOTAL_PASS}${RESET}"
    field_color "Warnings:" "${YELLOW}${TOTAL_WARN}${RESET}"
    field_color "Failures:" "${RED}${TOTAL_FAIL}${RESET}"
    field "Duration:" "$(elapsed)"
}

compute_exit() {
    case "$FAIL_ON" in
        warn)
            if [[ $TOTAL_FAIL -gt 0 || $TOTAL_WARN -gt 0 ]]; then
                exit 2
            fi
            ;;
        deny)
            if [[ $TOTAL_FAIL -gt 0 ]]; then
                exit 2
            fi
            ;;
    esac
}

# ── Output formats ───────────────────────────────────────────────────
write_output() {
    [[ -z "$OUTPUT_FILE" ]] && return
    case "$OUTPUT_FORMAT" in
        json) write_json ;;
        junit) write_junit ;;
        tap) write_tap ;;
        text) write_text ;;
        *) die "Unknown output format: ${OUTPUT_FORMAT}" ;;
    esac
    log "Results written to ${OUTPUT_FILE}"
}

write_json() {
    jq -n --argjson v "$RESULTS_JSON" \
        --argjson inputs "$TOTAL_INPUTS" \
        --argjson policies "$TOTAL_POLICIES" \
        --argjson pass "$TOTAL_PASS" \
        --argjson warns "$TOTAL_WARN" \
        --argjson fails "$TOTAL_FAIL" \
        '{summary: {inputs: $inputs, policies: $policies, pass: $pass, warnings: $warns, failures: $fails}, violations: $v}' \
        > "$OUTPUT_FILE"
}

write_junit() {
    local total tests failures
    total=$(echo "$RESULTS_JSON" | jq 'length')
    tests=$((TOTAL_PASS + total))
    failures=$TOTAL_FAIL

    {
        echo '<?xml version="1.0" encoding="UTF-8"?>'
        echo "<testsuites tests=\"${tests}\" failures=\"${failures}\">"
        echo "  <testsuite name=\"opa-policy-tester\" tests=\"${tests}\" failures=\"${failures}\">"

        local i=0
        while [[ $i -lt $TOTAL_PASS ]]; do
            echo "    <testcase name=\"input-pass-$((i+1))\" classname=\"opa\"/>"
            ((i++)) || true
        done

        echo "$RESULTS_JSON" | jq -r '.[] | @base64' | while read -r entry; do
            local vf vs vm
            vf=$(echo "$entry" | base64 -d | jq -r '.file')
            vs=$(echo "$entry" | base64 -d | jq -r '.severity')
            vm=$(echo "$entry" | base64 -d | jq -r '.message')
            echo "    <testcase name=\"${vf}:${vs}\" classname=\"opa\">"
            echo "      <failure message=\"${vm}\" type=\"${vs}\"/>"
            echo "    </testcase>"
        done

        echo "  </testsuite>"
        echo "</testsuites>"
    } > "$OUTPUT_FILE"
}

write_tap() {
    local total
    total=$(echo "$RESULTS_JSON" | jq 'length')
    local plan=$((TOTAL_PASS + total))
    {
        echo "TAP version 13"
        echo "1..${plan}"
        local idx=1
        local i=0
        while [[ $i -lt $TOTAL_PASS ]]; do
            echo "ok ${idx} - input passed all policies"
            ((idx++)) || true
            ((i++)) || true
        done

        echo "$RESULTS_JSON" | jq -r '.[] | [.file, .severity, .message] | @tsv' | \
        while IFS=$'\t' read -r vf vs vm; do
            echo "not ok ${idx} - ${vf}: [${vs}] ${vm}"
            ((idx++)) || true
        done
    } > "$OUTPUT_FILE"
}

write_text() {
    {
        echo "OPA Policy Test Results"
        echo "======================"
        echo ""
        echo "$RESULTS_JSON" | jq -r '.[] | "[\(.severity | ascii_upcase)] \(.file): \(.message)"'
        echo ""
        echo "Inputs: ${TOTAL_INPUTS}  Policies: ${TOTAL_POLICIES}  Pass: ${TOTAL_PASS}  Warn: ${TOTAL_WARN}  Fail: ${TOTAL_FAIL}"
    } > "$OUTPUT_FILE"
}

# ── Run unit tests ───────────────────────────────────────────────────
run_test() {
    [[ -z "$POLICY_DIR" ]] && die "No policy path specified (--policy)"
    [[ ! -e "$POLICY_DIR" ]] && die "Policy path does not exist: ${POLICY_DIR}"
    require_opa

    START_TIME=$(date +%s)
    section_header "OPA Unit Tests"
    field "Policy path:" "$POLICY_DIR"
    echo ""

    local test_output exit_code=0
    test_output=$("$OPA_PATH" test --verbose "$POLICY_DIR" 2>&1) || exit_code=$?

    local test_pass=0 test_fail=0
    while IFS= read -r line; do
        if [[ "$line" =~ ^PASS ]]; then
            ((test_pass++)) || true
            echo -e "  ${GREEN}✓${RESET} ${line#PASS }"
        elif [[ "$line" =~ ^FAIL ]]; then
            ((test_fail++)) || true
            echo -e "  ${RED}✗${RESET} ${line#FAIL }"
        elif [[ -n "$line" ]]; then
            verbose "$line"
        fi
    done <<< "$test_output"

    section_header "Test Summary"
    field_color "Passed:" "${GREEN}${test_pass}${RESET}"
    field_color "Failed:" "${RED}${test_fail}${RESET}"
    field "Duration:" "$(elapsed)"

    [[ $test_fail -gt 0 ]] && exit 2
    [[ $exit_code -ne 0 ]] && exit 1
}

# ── Format Rego files ────────────────────────────────────────────────
run_fmt() {
    [[ -z "$POLICY_DIR" ]] && die "No policy path specified (--policy)"
    [[ ! -e "$POLICY_DIR" ]] && die "Policy path does not exist: ${POLICY_DIR}"
    require_opa

    section_header "OPA Format"
    field "Policy path:" "$POLICY_DIR"
    echo ""

    local files=() formatted=0
    if [[ -f "$POLICY_DIR" ]]; then
        files=("$POLICY_DIR")
    else
        while IFS= read -r -d '' f; do
            files+=("$f")
        done < <(find "$POLICY_DIR" -name '*.rego' -print0 | sort -z)
    fi

    [[ ${#files[@]} -eq 0 ]] && die "No .rego files found in ${POLICY_DIR}"

    for f in "${files[@]}"; do
        local before after
        before=$(cat "$f")
        "$OPA_PATH" fmt -w "$f"
        after=$(cat "$f")
        if [[ "$before" != "$after" ]]; then
            ((formatted++)) || true
            echo -e "  ${YELLOW}reformatted${RESET} $(basename "$f")"
        else
            verbose "unchanged: $(basename "$f")"
        fi
    done

    echo ""
    field "Total files:" "${#files[@]}"
    field "Reformatted:" "$formatted"
}

# ── Syntax check ─────────────────────────────────────────────────────
run_check() {
    [[ -z "$POLICY_DIR" ]] && die "No policy path specified (--policy)"
    [[ ! -e "$POLICY_DIR" ]] && die "Policy path does not exist: ${POLICY_DIR}"
    require_opa

    section_header "OPA Syntax Check"
    field "Policy path:" "$POLICY_DIR"
    echo ""

    local check_output exit_code=0
    check_output=$("$OPA_PATH" check "$POLICY_DIR" 2>&1) || exit_code=$?

    if [[ $exit_code -eq 0 ]]; then
        echo -e "  ${GREEN}✓${RESET} All policies pass syntax check"
    else
        echo -e "  ${RED}✗${RESET} Syntax errors found:"
        echo ""
        while IFS= read -r line; do
            echo "    $line"
        done <<< "$check_output"
        exit 1
    fi
}

# ── Usage ─────────────────────────────────────────────────────────────
usage() {
    cat <<EOF
${BOLD}${SCRIPT_NAME}${RESET} — OPA Policy Tester

Run OPA/Rego policy bundles against infrastructure configs and report
violations with severity levels and CI-friendly exit codes.

${BOLD}MODES${RESET}
  --eval              Evaluate policies against input configs
  --test              Run OPA unit tests on policy directory
  --fmt               Format .rego files (opa fmt)
  --check             Syntax-check .rego files (opa check)
  --install           Install OPA binary if not present

${BOLD}OPTIONS${RESET}
  --policy DIR|FILE   Path to policy directory or .rego file
  --input  FILE|DIR   Input data file (JSON/YAML) or directory
  --data   FILE       Additional data file for evaluation
  --query  QUERY      Custom OPA deny query (default: data.main.deny)
  --fail-on LEVEL     Fail threshold: warn or deny (default: deny)
  --format  FORMAT    Output format: text, json, junit, tap (default: text)
  --output-file FILE  Write results to file
  --verbose           Show debug output
  --no-color          Disable colored output
  --help              Show this help message

${BOLD}ENVIRONMENT VARIABLES${RESET}
  OPA_PATH            Path to opa binary (default: opa)
  OPA_POLICY_DIR      Default policy directory
  OPA_INPUT           Default input path
  OPA_QUERY           Default deny query
  OPA_FORMAT          Default output format
  OPA_FAIL_ON         Default fail threshold
  VERBOSE             Enable verbose output (true/false)
  COLOR               Color mode: auto, always, never

${BOLD}EXAMPLES${RESET}
  # Evaluate Kubernetes manifests
  ${SCRIPT_NAME} --eval --policy ./policies --input ./k8s-manifests/

  # Check a single Terraform plan
  ${SCRIPT_NAME} --eval --policy ./policies --input plan.json --format json

  # Run unit tests
  ${SCRIPT_NAME} --test --policy ./policies

  # Syntax check and format
  ${SCRIPT_NAME} --check --policy ./policies
  ${SCRIPT_NAME} --fmt --policy ./policies

  # CI pipeline with strict threshold
  ${SCRIPT_NAME} --eval --policy ./policies --input . --fail-on warn --format junit --output-file results.xml

${BOLD}EXIT CODES${RESET}
  0   All checks passed
  1   Runtime error
  2   Violations exceed threshold
EOF
}

# ── Parse arguments ──────────────────────────────────────────────────
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --eval)     RUN_MODE="eval"; shift ;;
            --test)     RUN_MODE="test"; shift ;;
            --fmt)      RUN_MODE="fmt"; shift ;;
            --check)    RUN_MODE="check"; shift ;;
            --install)  RUN_MODE="install"; shift ;;
            --policy)   POLICY_DIR="${2:?--policy requires a path}"; shift 2 ;;
            --input)    INPUT_PATH="${2:?--input requires a path}"; shift 2 ;;
            --data)     DATA_FILE="${2:?--data requires a path}"; shift 2 ;;
            --query)    QUERY_DENY="${2:?--query requires a value}"; shift 2 ;;
            --fail-on)  FAIL_ON="${2:?--fail-on requires a value}"; shift 2 ;;
            --format)   OUTPUT_FORMAT="${2:?--format requires a value}"; shift 2 ;;
            --output-file) OUTPUT_FILE="${2:?--output-file requires a path}"; shift 2 ;;
            --verbose)  VERBOSE="true"; shift ;;
            --no-color) COLOR="never"; shift ;;
            --help|-h)  setup_colors; usage; exit 0 ;;
            *)          die "Unknown option: $1 (see --help)" ;;
        esac
    done
}

# ── Main ─────────────────────────────────────────────────────────────
main() {
    parse_args "$@"
    setup_colors

    case "$RUN_MODE" in
        eval)    run_eval ;;
        test)    run_test ;;
        fmt)     run_fmt ;;
        check)   run_check ;;
        install) install_opa ;;
        "")      err "No mode specified"; echo ""; usage; exit 1 ;;
        *)       die "Unknown mode: ${RUN_MODE}" ;;
    esac
}

main "$@"