#!/usr/bin/env bash
# =============================================================================
# fitness.sh — Push3 optimizer fitness scoring wrapper
#
# Pipeline: Push3 candidate → transpile → compile → deploy/upgrade → attack ×N → score
#
# Usage:
#   ./tools/push3-evolution/fitness.sh <candidate.push3>
#
# Output:
#   Single integer on stdout — total lm_eth_total across all attacks (wei).
#
# Exit codes:
#   0  Success — score printed to stdout.
#   1  Invalid candidate — Push3 program won't transpile, compile, or deploy.
#   2  Infra error — Anvil unavailable, missing tool, bootstrap failure.
#
# Environment:
#   ANVIL_FORK_URL   Required when Anvil is not already running.  Must point to
#                    a Base RPC endpoint so Uniswap V3 Factory and WETH exist at
#                    their canonical addresses (e.g. https://mainnet.base.org or
#                    a local Base fork).  Has no effect when Anvil is already up.
# =============================================================================

set -euo pipefail

# Foundry tools (forge, cast, anvil)
export PATH="${HOME}/.foundry/bin:${PATH}"

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
ONCHAIN_DIR="$REPO_ROOT/onchain"
ATTACKS_DIR="$ONCHAIN_DIR/script/backtesting/attacks"
RPC_URL="http://localhost:8545"

# Standard Anvil test accounts (deterministic mnemonic)
MNEMONIC="test test test test test test test test test test test junk"
# Account 2 — recenter caller
RECENTER_PK="0x5de4111afa1a4b94908f83103eb1f1706367c2e68ca870fc3fb9a804cdab365a"
# Account 8 — adversary (used to fund LM with WETH)
ADV_PK="0xdbda1821b80551c9d65939329250298aa3472ba22feea921c0cf5d620ea67b97"
# WETH address on the Base network
WETH="0x4200000000000000000000000000000000000006"

# =============================================================================
# Argument parsing
# =============================================================================

if [ $# -ne 1 ]; then
    echo "Usage: $0 <candidate.push3>" >&2
    exit 2
fi

PUSH3_FILE="$1"

if [ ! -f "$PUSH3_FILE" ]; then
    echo "Error: File not found: $PUSH3_FILE" >&2
    exit 1
fi

# Canonicalise so relative paths work after cwd changes.
PUSH3_FILE="$(cd "$(dirname "$PUSH3_FILE")" && pwd)/$(basename "$PUSH3_FILE")"

# =============================================================================
# Helpers
# =============================================================================

log()   { echo "  [fitness]  $*" >&2; }
fail1() { echo "  [invalid]  $*" >&2; exit 1; }
fail2() { echo "  [infra]    $*" >&2; exit 2; }

# =============================================================================
# Tool check
# =============================================================================

for _tool in forge cast anvil python3; do
    command -v "$_tool" &>/dev/null || fail2 "$_tool not found in PATH"
done

# =============================================================================
# Cleanup
#
# If we own Anvil (ANVIL_PID set), just kill it — no state cleanup needed.
# If we are using a shared Anvil (ANVIL_PID empty), revert to BASELINE_SNAP to
# undo bootstrap mutations (setRecenterAccess, WETH funding, initial recenter)
# so the chain is clean for the next caller.
# =============================================================================

ANVIL_PID=""
WORK_DIR="$(mktemp -d)"
BASELINE_SNAP=""   # pre-bootstrap snapshot; used to clean up on shared Anvil

cleanup() {
    if [ -n "$ANVIL_PID" ]; then
        kill "$ANVIL_PID" 2>/dev/null || true
    elif [ -n "$BASELINE_SNAP" ]; then
        cast rpc anvil_revert "$BASELINE_SNAP" --rpc-url "$RPC_URL" >/dev/null 2>&1 || true
    fi
    rm -rf "$WORK_DIR"
}
trap cleanup EXIT

# =============================================================================
# Step 0 — Ensure Anvil is running
#
# DeployLocal.sol depends on live Base infrastructure: Uniswap V3 Factory at
# 0x4752ba5DBc23f44D87826276BF6Fd6b1C372aD24 and WETH at
# 0x4200000000000000000000000000000000000006.  A plain (unfork'd) Anvil has
# neither, so cold-starting without --fork-url silently breaks the pipeline.
#
# When Anvil is already running (dev stack or CI), we use it as-is.
# When it is not running we require ANVIL_FORK_URL and start a forked instance.
# =============================================================================

if cast chain-id --rpc-url "$RPC_URL" >/dev/null 2>&1; then
    log "Anvil already running at $RPC_URL"
else
    ANVIL_FORK_URL="${ANVIL_FORK_URL:-}"
    if [ -z "$ANVIL_FORK_URL" ]; then
        fail2 "Anvil is not running at $RPC_URL and ANVIL_FORK_URL is not set.
  DeployLocal.sol requires Base network contracts (Uniswap V3 Factory, WETH).
  Either start a Base-forked Anvil externally, or set ANVIL_FORK_URL to a Base
  RPC endpoint (e.g. ANVIL_FORK_URL=https://mainnet.base.org)."
    fi

    anvil --silent \
          --fork-url "$ANVIL_FORK_URL" \
          --mnemonic "$MNEMONIC" \
          --port 8545 &
    ANVIL_PID=$!

    TRIES=0
    until cast chain-id --rpc-url "$RPC_URL" >/dev/null 2>&1; do
        TRIES=$((TRIES + 1))
        [ $TRIES -gt 50 ] && fail2 "Anvil did not start within 50 attempts"
        sleep 0.2
    done
    log "Anvil started (PID $ANVIL_PID, fork: $ANVIL_FORK_URL)"
fi

# =============================================================================
# Steps 1–3 — Transpile → compile → deploy fresh stack → UUPS upgrade
#
# BASELINE_SNAP is taken before deploy-optimizer.sh runs so that cleanup()
# can fully revert the deploy on a shared Anvil (not just the bootstrap
# mutations).  Without this, a second sequential evaluation against the same
# shared Anvil would run DeployLocal.sol again with the same deployer nonce,
# hitting CREATE address collisions and failing with exit 1.
#
# deploy-optimizer.sh handles the full pipeline.  With no OPTIMIZER_PROXY set
# it also runs DeployLocal.sol to produce the initial stack.
#
# Output is captured silently on success; surfaced to stderr on failure so
# batch / evolution-loop callers are not flooded with per-candidate progress
# lines for every successful evaluation.
#
# Exit codes from deploy-optimizer.sh all map to exit 1 (invalid candidate)
# because transpile / compile / round-trip failures are candidate issues.
# =============================================================================

# Pre-deploy snapshot — reverted in cleanup to fully undo the deploy and all
# bootstrap mutations when using a shared Anvil.
BASELINE_SNAP=$(cast rpc anvil_snapshot --rpc-url "$RPC_URL" | tr -d '"')
log "Pre-deploy snapshot: $BASELINE_SNAP"

log "Running deploy-optimizer.sh (transpile → compile → deploy → upgrade)…"

DEPLOY_LOG="$WORK_DIR/deploy.log"
DEPLOY_EC=0
"$REPO_ROOT/tools/deploy-optimizer.sh" "$PUSH3_FILE" >"$DEPLOY_LOG" 2>&1 || DEPLOY_EC=$?

if [ "$DEPLOY_EC" -ne 0 ]; then
    # Surface deploy log so operators can diagnose candidate failures.
    cat "$DEPLOY_LOG" >&2
    fail1 "deploy-optimizer.sh failed (exit $DEPLOY_EC)"
fi

log "Optimizer deployed and upgraded"

# =============================================================================
# Step 4 — Read deployment addresses
#
# DeployLocal.sol writes deterministic addresses to deployments-local.json when
# run against a fresh Anvil + standard mnemonic.
# =============================================================================

DEPLOYMENTS="$ONCHAIN_DIR/deployments-local.json"
[ -f "$DEPLOYMENTS" ] || fail2 "deployments-local.json not found — did DeployLocal.sol run?"

LM_ADDR=$(python3 -c "
import json
d = json.load(open('$DEPLOYMENTS'))
print(d['contracts']['LiquidityManager'])
" 2>/dev/null) || fail2 "Failed to read LiquidityManager from deployments-local.json"

[ -n "$LM_ADDR" ] || fail2 "LiquidityManager address is empty in deployments-local.json"
log "LiquidityManager: $LM_ADDR"

# =============================================================================
# Step 5 — Bootstrap LM state
#
# a. Grant recenterAccess to account 2 (impersonate feeDestination).
# b. Fund LM with 1000 WETH from account 8.
# c. Call recenter() to deploy capital into Uniswap positions.
#    The LM needs TWAP history; mine blocks in batches and retry.
# =============================================================================

# a. Grant recenterAccess.
RECENTER_ADDR=$(cast wallet address --private-key "$RECENTER_PK")

FEE_DEST=$(cast call "$LM_ADDR" "feeDestination()(address)" \
    --rpc-url "$RPC_URL" 2>/dev/null | sed 's/\[.*//;s/[[:space:]]//g') \
    || fail2 "Failed to read feeDestination() from LM"

log "Granting recenterAccess to $RECENTER_ADDR (via feeDestination $FEE_DEST)"
cast rpc --rpc-url "$RPC_URL" anvil_impersonateAccount "$FEE_DEST" >/dev/null
cast send --rpc-url "$RPC_URL" --from "$FEE_DEST" --unlocked \
    "$LM_ADDR" "setRecenterAccess(address)" "$RECENTER_ADDR" >/dev/null 2>&1 \
    || fail2 "setRecenterAccess failed"
cast rpc --rpc-url "$RPC_URL" anvil_stopImpersonatingAccount "$FEE_DEST" >/dev/null

# c. Fund LM with 1000 WETH.
log "Funding LM with 1000 WETH"
cast send "$WETH" "deposit()" --value 1000ether \
    --private-key "$ADV_PK" --rpc-url "$RPC_URL" >/dev/null 2>&1 \
    || fail2 "Failed to wrap ETH to WETH"
cast send "$WETH" "transfer(address,uint256)" "$LM_ADDR" 1000000000000000000000 \
    --private-key "$ADV_PK" --rpc-url "$RPC_URL" >/dev/null 2>&1 \
    || fail2 "Failed to transfer WETH to LM"

# d. Initial recenter.  Mine 50 blocks per attempt (single anvil_mine call) to
#    build the TWAP history the LM needs before recenter() will succeed.
log "Initial recenter — deploying capital into positions"
RECENTERED=false
for _attempt in 1 2 3 4; do
    cast rpc --rpc-url "$RPC_URL" anvil_mine 0x32 >/dev/null 2>&1
    if cast send "$LM_ADDR" "recenter()" \
            --private-key "$RECENTER_PK" --rpc-url "$RPC_URL" >/dev/null 2>&1; then
        RECENTERED=true
        break
    fi
done

if ! $RECENTERED; then
    log "WARNING: initial recenter did not succeed — attack scores may be lower than expected"
fi

# =============================================================================
# Steps 6–7 — Run each attack and accumulate lm_eth_total
#
# Each attack starts from the same post-bootstrap state by taking a snapshot
# before the run and reverting to it afterwards.  If a revert fails (wrong
# snapshot ID, Anvil restart, etc.) subsequent attacks would run against dirty
# state; we flag this so the caller can discard the score.
# =============================================================================

TOTAL_ETH=0
ATTACK_COUNT=0
DIRTY=false

for ATTACK_JSONL in "$ATTACKS_DIR"/*.jsonl; do
    [ -f "$ATTACK_JSONL" ] || continue
    ATTACK_NAME="$(basename "$ATTACK_JSONL" .jsonl)"
    log "Running attack: $ATTACK_NAME"

    # a. Take per-attack snapshot.
    ATK_SNAP=$(cast rpc anvil_snapshot --rpc-url "$RPC_URL" | tr -d '"')

    # b. Run AttackRunner, capturing all output.
    #    console.log snapshot lines start with '{'; other forge output does not.
    ATK_OUT="$WORK_DIR/atk-${ATTACK_NAME}.txt"
    ATK_EC=0
    (
        cd "$ONCHAIN_DIR"
        ATTACK_FILE="$ATTACK_JSONL" \
        forge script script/backtesting/AttackRunner.s.sol \
            --rpc-url "$RPC_URL" --broadcast --no-color 2>&1
    ) >"$ATK_OUT" || ATK_EC=$?

    if [ "$ATK_EC" -ne 0 ]; then
        log "  WARNING: AttackRunner failed for $ATTACK_NAME (exit $ATK_EC) — skipping"
        if ! cast rpc anvil_revert "$ATK_SNAP" --rpc-url "$RPC_URL" >/dev/null 2>&1; then
            log "  WARNING: anvil_revert also failed — subsequent attacks run against dirty state"
            DIRTY=true
        fi
        continue
    fi

    # c. Extract lm_eth_total from the final JSON snapshot line.
    ETH_RETAINED=$(python3 - "$ATK_OUT" <<'PYEOF'
import sys, json
snapshots = []
with open(sys.argv[1]) as f:
    for line in f:
        line = line.strip()
        if line.startswith('{') and '"lm_eth_total"' in line:
            try:
                snapshots.append(json.loads(line))
            except json.JSONDecodeError:
                pass
if snapshots:
    # lm_eth_total is serialised as a quoted integer string.
    val = snapshots[-1]['lm_eth_total']
    print(int(val) if isinstance(val, str) else val)
else:
    print(0)
PYEOF
)

    log "  $ATTACK_NAME: lm_eth_total=$ETH_RETAINED"
    TOTAL_ETH=$(python3 -c "print(int('$TOTAL_ETH') + int('$ETH_RETAINED'))")
    ATTACK_COUNT=$((ATTACK_COUNT + 1))

    # d. Revert to per-attack snapshot to restore post-bootstrap baseline.
    if ! cast rpc anvil_revert "$ATK_SNAP" --rpc-url "$RPC_URL" >/dev/null 2>&1; then
        log "  WARNING: anvil_revert failed for $ATTACK_NAME — subsequent attacks run against dirty state"
        DIRTY=true
    fi
done

# =============================================================================
# Output
# =============================================================================

if [ "$ATTACK_COUNT" -eq 0 ]; then
    fail2 "No attacks ran — check $ATTACKS_DIR for *.jsonl files"
fi

if $DIRTY; then
    log "WARNING: one or more revert failures occurred — score may be inaccurate"
fi

log "Score: $TOTAL_ETH wei (sum of lm_eth_total across $ATTACK_COUNT attacks)"
echo "$TOTAL_ETH"