#!/usr/bin/env bash
# Ergonomic wrapper for running the eval pipeline inside a Docker container.
#
# Usage:
#   evals/docker/run_eval_docker.sh [run_eval.py args...]
#
# Examples:
#   # Smoke test (no network needed)
#   evals/docker/run_eval_docker.sh --patcher null --max-cases 5 --k 1
#
#   # OpenAI patcher against local vLLM
#   evals/docker/run_eval_docker.sh \
#     --patcher openai --openai-model mistral-7b \
#     --openai-base-url http://127.0.0.1:8000/v1 \
#     --category hidden_state --k 3
#
#   # Run from TOML config (paths are staged into the container automatically)
#   evals/docker/run_eval_docker.sh \
#     --patcher openai \
#     --config evals/configs/openai_adaptive.template.toml
#
#   # Force rebuild
#   LLMDEBUG_EVAL_BUILD=1 evals/docker/run_eval_docker.sh --patcher null --max-cases 1
#
# Environment variables:
#   LLMDEBUG_EVAL_BUILD=1          Force image rebuild
#   LLMDEBUG_EVAL_IMAGE=name:tag   Override image name (default: llmdebug-eval:latest)
#   DOCKER_EXTRA_ARGS="..."        Extra docker run flags (e.g. --gpus all)
#   OPENAI_MODEL=...               Default model when --patcher openai is used
#   OPENAI_BASE_URL=...            Default base URL when --patcher openai is used
#   Z_API / OPENAI_API_KEY / LLMDEBUG_EVAL_OPENAI_API_KEY
#                                  Used to synthesize --openai-api-key when omitted
set -euo pipefail

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

IMAGE="${LLMDEBUG_EVAL_IMAGE:-llmdebug-eval:latest}"

# ---------------------------------------------------------------------------
# Auto-build if image doesn't exist or rebuild requested
# ---------------------------------------------------------------------------
if [[ "${LLMDEBUG_EVAL_BUILD:-}" == "1" ]] || ! docker image inspect "$IMAGE" &>/dev/null; then
    echo ">>> Building $IMAGE ..."
    docker build -t "$IMAGE" -f "$REPO_ROOT/evals/docker/eval.Dockerfile" "$REPO_ROOT"
    echo ">>> Build complete."
fi

# ---------------------------------------------------------------------------
# Prepare docker run arguments
# ---------------------------------------------------------------------------
DOCKER_ARGS=(
    --rm
    -v "$REPO_ROOT/evals/results:/repo/evals/results"
    -v "$REPO_ROOT/evals/artifacts:/repo/evals/artifacts"
)
DEFAULT_DOCKER_USER="$(id -u):$(id -g)"
DOCKER_USER="${LLMDEBUG_EVAL_DOCKER_USER:-$DEFAULT_DOCKER_USER}"

# Forward non-secret config env vars if set.
# OpenAI credentials are passed via CLI flags (only when needed) so that
# benchmarked test code cannot read them from the environment.
for var in LLMDEBUG_EVAL_PATCH_CMD LLMDEBUG_EVAL_PATCH_CMD_TIMEOUT_SEC \
           LLMDEBUG_EVAL_PATCH_CMD_SHELL HF_HOME; do
    if [[ -n "${!var:-}" ]]; then
        DOCKER_ARGS+=(-e "$var=${!var}")
    fi
done

# Extra user-specified docker flags
if [[ -n "${DOCKER_EXTRA_ARGS:-}" ]]; then
    read -ra _extra <<< "$DOCKER_EXTRA_ARGS"
    DOCKER_ARGS+=("${_extra[@]}")
fi

# ---------------------------------------------------------------------------
# Argument helpers
# ---------------------------------------------------------------------------
_arg_value() {
    local key="$1"
    shift
    local next_is_value=0
    local found=0
    local value=""
    local arg
    for arg in "$@"; do
        if [[ $next_is_value -eq 1 ]]; then
            value="$arg"
            found=1
            next_is_value=0
            continue
        fi
        if [[ "$arg" == "$key" ]]; then
            next_is_value=1
            continue
        fi
        if [[ "$arg" == "$key="* ]]; then
            value="${arg#*=}"
            found=1
        fi
    done
    if [[ $found -eq 1 ]]; then
        printf "%s\n" "$value"
        return 0
    fi
    return 1
}

_arg_present() {
    local key="$1"
    shift
    local arg
    for arg in "$@"; do
        if [[ "$arg" == "$key" || "$arg" == "$key="* ]]; then
            return 0
        fi
    done
    return 1
}

_replace_or_append_arg() {
    local key="$1"
    local value="$2"
    local out=()
    local replaced=0
    local skip_next=0
    local arg

    for arg in "${EVAL_ARGS[@]}"; do
        if [[ $skip_next -eq 1 ]]; then
            skip_next=0
            continue
        fi
        if [[ "$arg" == "$key" ]]; then
            out+=("$key" "$value")
            replaced=1
            skip_next=1
            continue
        fi
        if [[ "$arg" == "$key="* ]]; then
            out+=("$key=$value")
            replaced=1
            continue
        fi
        out+=("$arg")
    done

    if [[ $replaced -eq 0 ]]; then
        out+=("$key" "$value")
    fi

    EVAL_ARGS=("${out[@]}")
}

_first_nonempty_env() {
    local var
    local value
    for var in "$@"; do
        value="${!var:-}"
        if [[ -n "$value" ]]; then
            printf "%s\n" "$value"
            return 0
        fi
    done
    return 1
}

_expand_config_path() {
    local raw="$1"
    local expanded="$raw"
    if [[ "$expanded" == "~/"* ]]; then
        expanded="$HOME/${expanded#~/}"
    fi
    if [[ "$expanded" != /* ]]; then
        expanded="$PWD/$expanded"
    fi
    printf "%s\n" "$expanded"
}

_rewrite_config_paths_for_container() {
    local config_stage_id
    config_stage_id="$(date +%Y%m%dT%H%M%S)-$$-$RANDOM"
    local host_mount_dir="$REPO_ROOT/evals/artifacts/config_mounts/$config_stage_id"
    local guest_mount_dir="/repo/evals/artifacts/config_mounts/$config_stage_id"
    local out=()
    local expect_value=0
    local config_index=0
    local arg

    for arg in "${EVAL_ARGS[@]}"; do
        if [[ $expect_value -eq 1 ]]; then
            local host_path
            host_path="$(_expand_config_path "$arg")"
            if [[ ! -f "$host_path" ]]; then
                echo ">>> ERROR: Config file not found: $arg" >&2
                exit 2
            fi
            mkdir -p "$host_mount_dir"
            config_index=$((config_index + 1))
            local staged_name
            staged_name="$(printf "%03d_%s" "$config_index" "$(basename "$host_path")")"
            cp "$host_path" "$host_mount_dir/$staged_name"
            out+=(--config "$guest_mount_dir/$staged_name")
            expect_value=0
            continue
        fi

        if [[ "$arg" == "--config" ]]; then
            expect_value=1
            continue
        fi
        if [[ "$arg" == --config=* ]]; then
            local raw_path="${arg#--config=}"
            local host_path
            host_path="$(_expand_config_path "$raw_path")"
            if [[ ! -f "$host_path" ]]; then
                echo ">>> ERROR: Config file not found: $raw_path" >&2
                exit 2
            fi
            mkdir -p "$host_mount_dir"
            config_index=$((config_index + 1))
            local staged_name
            staged_name="$(printf "%03d_%s" "$config_index" "$(basename "$host_path")")"
            cp "$host_path" "$host_mount_dir/$staged_name"
            out+=("--config=$guest_mount_dir/$staged_name")
            continue
        fi

        out+=("$arg")
    done

    if [[ $expect_value -eq 1 ]]; then
        echo ">>> ERROR: --config requires a path value" >&2
        exit 2
    fi

    EVAL_ARGS=("${out[@]}")
}

_docker_has_network_flag() {
    local expect_value=0
    local arg
    for arg in "${DOCKER_ARGS[@]}"; do
        if [[ $expect_value -eq 1 ]]; then
            return 0
        fi
        case "$arg" in
            --network|--net)
                expect_value=1
                ;;
            --network=*|--net=*)
                return 0
                ;;
        esac
    done
    return 1
}

_docker_user_can_write_outputs() {
    local user="$1"
    docker run \
        --rm \
        --entrypoint /bin/sh \
        --user "$user" \
        -v "$REPO_ROOT/evals/results:/repo/evals/results" \
        -v "$REPO_ROOT/evals/artifacts:/repo/evals/artifacts" \
        "$IMAGE" \
        -lc 'touch /repo/evals/results/.perm_probe && rm -f /repo/evals/results/.perm_probe && touch /repo/evals/artifacts/.perm_probe && rm -f /repo/evals/artifacts/.perm_probe' \
        >/dev/null 2>&1
}

# ---------------------------------------------------------------------------
# Filter and forward arguments to evals.run_eval
# ---------------------------------------------------------------------------
EVAL_ARGS=()
skip_next=0
for arg in "$@"; do
    if [[ $skip_next -eq 1 ]]; then
        skip_next=0
        continue
    fi

    case "$arg" in
        --executor)
            skip_next=1
            ;;
        --executor=*)
            ;;
        --container-image|--container-network|--container-cpus|--container-memory|\
        --container-pids-limit|--container-user|--container-tmpfs-size-mb|--container-env)
            skip_next=1
            ;;
        --container-image=*|--container-network=*|--container-cpus=*|--container-memory=*|\
        --container-pids-limit=*|--container-user=*|--container-tmpfs-size-mb=*|--container-env=*)
            ;;
        --container-mount-repo-src|--container-no-mount-repo-src)
            ;;
        --container-*)
            echo ">>> NOTE: Ignoring unsupported wrapper-only arg: $arg" >&2
            ;;
        *)
            EVAL_ARGS+=("$arg")
            ;;
    esac
done

# Stage and rewrite --config file paths so they are readable inside the container.
_rewrite_config_paths_for_container

# Resolve patcher mode (defaults to command when omitted)
PATCHER="$(_arg_value --patcher "${EVAL_ARGS[@]}" || true)"
PATCHER="${PATCHER:-command}"

# OpenAI defaults are resolved into CLI args so values are deterministic and
# localhost URL translation can safely rewrite the effective value.
if [[ "$PATCHER" == "openai" ]]; then
    if ! _arg_present --openai-model "${EVAL_ARGS[@]}" && [[ -n "${OPENAI_MODEL:-}" ]]; then
        EVAL_ARGS+=(--openai-model "$OPENAI_MODEL")
    fi

    if ! _arg_present --openai-base-url "${EVAL_ARGS[@]}"; then
        EVAL_ARGS+=(--openai-base-url "${OPENAI_BASE_URL:-http://127.0.0.1:1234}")
    fi

    if ! _arg_present --openai-api-key "${EVAL_ARGS[@]}"; then
        OPENAI_API_KEY_EFFECTIVE="$(_first_nonempty_env Z_API OPENAI_API_KEY LLMDEBUG_EVAL_OPENAI_API_KEY || true)"
        if [[ -n "${OPENAI_API_KEY_EFFECTIVE:-}" ]]; then
            EVAL_ARGS+=(--openai-api-key "$OPENAI_API_KEY_EFFECTIVE")
        fi
    fi
fi

# Localhost API URL handling for Docker networking
OPENAI_BASE_URL_EFFECTIVE="$(_arg_value --openai-base-url "${EVAL_ARGS[@]}" || true)"
if [[ -n "${OPENAI_BASE_URL_EFFECTIVE:-}" ]] \
    && printf "%s" "$OPENAI_BASE_URL_EFFECTIVE" | grep -qE '(127\.0\.0\.1|localhost)'; then
    if [[ "$(uname -s)" == "Darwin" ]]; then
        NEW_OPENAI_BASE_URL="$(printf "%s" "$OPENAI_BASE_URL_EFFECTIVE" | sed -E 's/(127\.0\.0\.1|localhost)/host.docker.internal/g')"
        if [[ "$NEW_OPENAI_BASE_URL" != "$OPENAI_BASE_URL_EFFECTIVE" ]]; then
            echo ">>> NOTE: Translating localhost URL for Docker access" >&2
            echo ">>>   $OPENAI_BASE_URL_EFFECTIVE -> $NEW_OPENAI_BASE_URL" >&2
            _replace_or_append_arg --openai-base-url "$NEW_OPENAI_BASE_URL"
        fi
    elif ! _docker_has_network_flag; then
        echo ">>> NOTE: Detected localhost API URL, using --network host" >&2
        DOCKER_ARGS+=(--network host)
    fi
fi

# Force local executor (the whole point of this container)
EVAL_ARGS+=(--executor local)

# ---------------------------------------------------------------------------
# Ensure output directories exist on host (docker mount creates them as root otherwise)
# ---------------------------------------------------------------------------
mkdir -p "$REPO_ROOT/evals/results" "$REPO_ROOT/evals/artifacts"

# On some Docker Desktop mount setups, host-owned dirs are presented as
# root-owned 755 inside the container. In that case, forcing --user UID:GID
# causes PermissionError when creating run artifact folders.
if [[ -n "$DOCKER_USER" ]]; then
    if [[ -n "${LLMDEBUG_EVAL_DOCKER_USER:-}" ]]; then
        DOCKER_ARGS+=(--user "$DOCKER_USER")
    elif _docker_user_can_write_outputs "$DOCKER_USER"; then
        DOCKER_ARGS+=(--user "$DOCKER_USER")
    else
        echo ">>> WARNING: Docker user $DOCKER_USER cannot write mounted eval outputs; falling back to root." >&2
        echo ">>>          Set LLMDEBUG_EVAL_DOCKER_USER to override this behavior." >&2
    fi
fi

# ---------------------------------------------------------------------------
# Run — exec replaces this shell so Ctrl-C propagates directly
# ---------------------------------------------------------------------------
exec docker run \
    "${DOCKER_ARGS[@]}" \
    --entrypoint python \
    "$IMAGE" \
    -m evals.run_eval \
    "${EVAL_ARGS[@]}"
