set -euo pipefail

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

if command -v uv &> /dev/null; then
    UV_BIN=$(command -v uv)
else
    if command -v python &> /dev/null; then
        PYTHON=python
    elif command -v python3 &> /dev/null; then
        PYTHON=python3
    else
        echo "Error: No Python interpreter found" >&2
        exit 1
    fi
    $PYTHON -m pip install uv &> /dev/null || true
    UV_BIN=$($PYTHON -c 'import uv; print(uv.find_uv_bin())' 2>/dev/null || command -v uv 2>/dev/null || echo "uv")
fi

# try SCRATCH -> TMPDIR -> /tmp for uv cache to avoid filesystem locking issues
# GROUNDHOG_CACHE_DIR allows users to override if needed, UV_* vars also respected
{% raw %}
GROUNDHOG_CACHE_BASE="${{GROUNDHOG_CACHE_DIR:-${{SCRATCH:-${{TMPDIR:-/tmp}}}}}}"

if [ -z "${{UV_CACHE_DIR:-}}" ]; then
    export UV_CACHE_DIR="${{GROUNDHOG_CACHE_BASE}}/${{USER:-$(id -un)}}/uv"
fi

if [ -z "${{UV_PYTHON_INSTALL_DIR:-}}" ]; then
    export UV_PYTHON_INSTALL_DIR="${{GROUNDHOG_CACHE_BASE}}/${{USER:-$(id -un)}}/uv/python"
fi
mkdir -p "$UV_CACHE_DIR" "$UV_PYTHON_INSTALL_DIR"
{% endraw %}

# Environment reuse: compute path from hash + version
ENV_HASH="{{ env_hash }}"
GROUNDHOG_VERSION="{{ groundhog_version }}"
{% raw %}
ENV_DIR="${{GROUNDHOG_CACHE_BASE}}/${{USER:-$(id -un)}}/groundhog-envs/${{ENV_HASH}}-${{GROUNDHOG_VERSION}}"
{% endraw %}

# Propagate log level to remote environment
{% if log_level %}
# Local override - use value from dispatching environment
export GROUNDHOG_LOG_LEVEL="{{ log_level }}"
{% else %}
{% raw %}
# Respect remote environment if set, otherwise default to WARNING
export GROUNDHOG_LOG_LEVEL="${{GROUNDHOG_LOG_LEVEL:-WARNING}}"
{% endraw %}
{% endif %}

cat > "$TASK_DIR/user_script.py" << 'USER_SCRIPT_EOF'
{{ user_script_contents | escape_braces }}
USER_SCRIPT_EOF

cat > "$TASK_DIR/runner.py" << 'RUNNER_EOF'
{{ runner_contents | escape_braces }}
RUNNER_EOF

cat > "$TASK_DIR/payload.in" << 'PAYLOAD_EOF'
{payload}
PAYLOAD_EOF

# Check if environment exists; create if not
if [ -d "$ENV_DIR" ]; then
    # Environment exists - reuse it
    {% raw %}
    if [ "${{GROUNDHOG_LOG_LEVEL}}" = "INFO" ] || [ "${{GROUNDHOG_LOG_LEVEL}}" = "DEBUG" ]; then
        echo "INFO: Using environment $ENV_DIR" >&2
    fi
    {% endraw %}
else
    # Create new environment
    {% raw %}
    if [ "${{GROUNDHOG_LOG_LEVEL}}" = "INFO" ] || [ "${{GROUNDHOG_LOG_LEVEL}}" = "DEBUG" ]; then
        echo "INFO: Creating environment $ENV_DIR" >&2
    fi
    {% endraw %}

    {% if uv_config_toml %}

    # Both `uv venv` and `uv pip install` will receive --config-file, then
    # the file is moved into $ENV_DIR as a human-readable audit trail.
    {% raw %}UV_TOML="${{ENV_DIR}}.uv.toml"{% endraw %}
    mkdir -p "$(dirname "$UV_TOML")"

    cat > "$UV_TOML" << 'UV_CONFIG_EOF'
{{ uv_config_toml | escape_braces }}
UV_CONFIG_EOF
    {% endif %}

    "$UV_BIN" venv "$ENV_DIR"{% if requires_python %} --python "{{ requires_python }}"{% endif %}{% if uv_config_toml %} --config-file "$UV_TOML"{% endif %}

    {% if uv_config_toml %}
    mv "$UV_TOML" "$ENV_DIR/uv.toml"
    {% endif %}

    # Install dependencies
    "$UV_BIN" pip install --python "$ENV_DIR/bin/python" \
        {% if uv_config_toml %}--config-file "$ENV_DIR/uv.toml" {% endif %}\
        --exclude-newer-package groundhog-hpc={{ groundhog_timestamp }} \
        {% for dep in dependencies %}"{{ dep }}" {% endfor %}{{ version_spec }}

    # Write metadata for posterity
    cat > "$ENV_DIR/groundhog-meta.json" << 'META_EOF'
{{ '{{' }}
    "created_at": "{{ groundhog_timestamp }}",
    "requires_python": "{{ requires_python }}",
    "dependencies": {{ dependencies | tojson }},
    "groundhog_version": "{{ groundhog_version }}"
{{ '}}' }}
META_EOF
fi

# Run using the cached environment's Python directly (bypasses uv resolution)
cd "$TASK_DIR"
"$ENV_DIR/bin/python" runner.py

echo "__GROUNDHOG_RESULT__"
cat payload.out
