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

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

# shellcheck source=/home/intlc/projects/proxmox/scripts/lib/load-project-env.sh
source "${PROJECT_ROOT}/scripts/lib/load-project-env.sh"

: "${ETHEREUM_MAINNET_RPC:?ETHEREUM_MAINNET_RPC is required}"

VAULT="${ENGINE_X_VAULT:-0x9a22a3e272A364D64240dE6bda796FcA421cA7E9}"
CWUSDC="${CWUSDC_MAINNET:-0x2de5F116bFcE3d0f922d9C8351e0c5Fc24b9284a}"
XAUT="${XAUT_MAINNET:-0x68749665FF8D2d112Fa859AA293F07A622782F38}"
LOOPS="${LOOPS:-1}"
EXACT_OUTPUT_PER_LOOP_RAW="${EXACT_OUTPUT_PER_LOOP_RAW:-10000}"
DEBT_RAW="${DEBT_RAW:-}"
ROUNDING_RECEIVER="${ROUNDING_RECEIVER:-}"
RECIPIENT="${RECIPIENT:-}"
EXECUTE="${EXECUTE:-0}"
STAMP="${ENGINE_X_LOOP_PROOF_STAMP:-$(date -u +%Y%m%dT%H%M%SZ)}"
OUT_JSON="${OUT_JSON:-reports/status/engine-x-loop-proof-${STAMP}.json}"
OUT_MD="${OUT_MD:-reports/status/engine-x-loop-proof-${STAMP}.md}"
LATEST_JSON="${LATEST_JSON:-reports/status/engine-x-loop-proof-latest.json}"
LATEST_MD="${LATEST_MD:-reports/status/engine-x-loop-proof-latest.md}"

if [[ -n "${PRIVATE_KEY:-}" ]]; then
  SIGNER="$(cast wallet address --private-key "${PRIVATE_KEY}")"
else
  SIGNER="${DEPLOYER_ADDRESS:-}"
fi
if [[ -z "${SIGNER}" ]]; then
  echo "Set PRIVATE_KEY or DEPLOYER_ADDRESS" >&2
  exit 1
fi
if [[ "${EXECUTE}" == "1" && -z "${PRIVATE_KEY:-}" ]]; then
  echo "PRIVATE_KEY is required when EXECUTE=1" >&2
  exit 1
fi

RECIPIENT="${RECIPIENT:-${SIGNER}}"
ROUNDING_RECEIVER="${ROUNDING_RECEIVER:-${SIGNER}}"
PROOF_ID="${PROOF_ID:-$(cast keccak "dbis-engine-x:loop-proof:${SIGNER}:${RECIPIENT}:${EXACT_OUTPUT_PER_LOOP_RAW}:${LOOPS}:${STAMP}")}"
ISO_HASH="${ISO_HASH:-$(cast keccak "dbis-engine-x:loop-proof:iso:${PROOF_ID}")}"
AUDIT_HASH="${AUDIT_HASH:-$(cast keccak "dbis-engine-x:loop-proof:audit:${PROOF_ID}")}"
PEG_HASH="${PEG_HASH:-$(cast keccak "dbis-engine-x:loop-proof:peg:${PROOF_ID}")}"

POOL_CWUSDC_RAW="$(cast call "${VAULT}" 'poolCwusdcReserve()(uint256)' --rpc-url "${ETHEREUM_MAINNET_RPC}" | awk '{print $1}')"
POOL_USDC_RAW="$(cast call "${VAULT}" 'poolUsdcReserve()(uint256)' --rpc-url "${ETHEREUM_MAINNET_RPC}" | awk '{print $1}')"
LENDER_USDC_RAW="$(cast call "${VAULT}" 'lenderUsdcAvailable()(uint256)' --rpc-url "${ETHEREUM_MAINNET_RPC}" | awk '{print $1}')"

if [[ -z "${DEBT_RAW}" ]]; then
  DEBT_RAW="$(
    python3 - "${POOL_CWUSDC_RAW}" "${POOL_USDC_RAW}" "${EXACT_OUTPUT_PER_LOOP_RAW}" "${LENDER_USDC_RAW}" <<'PY'
import sys

reserve_cw = int(sys.argv[1])
reserve_usdc = int(sys.argv[2])
target = int(sys.argv[3])
max_debt = int(sys.argv[4])

def amount_in(amount_out, reserve_in, reserve_out):
    if amount_out <= 0 or reserve_in <= 0 or reserve_out <= amount_out:
        raise ValueError("bad getAmountIn inputs")
    return reserve_in * amount_out * 1000 // ((reserve_out - amount_out) * 997) + 1

def amount_out(amount_in, reserve_in, reserve_out):
    if amount_in <= 0 or reserve_in <= 0 or reserve_out <= 0:
        raise ValueError("bad getAmountOut inputs")
    return amount_in * 997 * reserve_out // (reserve_in * 1000 + amount_in * 997)

for debt in range(1, max_debt + 1):
    cw_in = amount_in(debt, reserve_cw, reserve_usdc)
    cw_reserve_after_in = reserve_cw + cw_in
    usdc_reserve_after_out = reserve_usdc - debt
    cw_out = amount_out(debt, usdc_reserve_after_out, cw_reserve_after_in)
    if cw_out >= target:
        print(debt)
        raise SystemExit(0)
raise SystemExit("no debt size can reach exact output within lender bucket")
PY
  )"
fi

mapfile -t PREVIEW < <(
  cast call "${VAULT}" \
    'previewVirtualProof(uint256,uint256)(uint256,uint256,uint256,uint256,uint256,uint256,uint256)' \
    "${DEBT_RAW}" "${LOOPS}" --rpc-url "${ETHEREUM_MAINNET_RPC}" | awk '{print $1}'
)

COLLATERAL_XAUT_RAW="${PREVIEW[0]}"
CWUSDC_IN_PER_LOOP_RAW="${PREVIEW[1]}"
CWUSDC_OUT_PER_LOOP_RAW="${PREVIEW[2]}"
CWUSDC_LOSS_PER_LOOP_RAW="${PREVIEW[3]}"
TOTAL_CWUSDC_IN_RAW="${PREVIEW[4]}"
TOTAL_CWUSDC_OUT_RAW="${PREVIEW[5]}"
TOTAL_NEUTRALIZED_RAW="${PREVIEW[6]}"
EXACT_OUTPUT_TOTAL_RAW="$((EXACT_OUTPUT_PER_LOOP_RAW * LOOPS))"
OUTPUT_ROUNDING_RAW="$((TOTAL_CWUSDC_OUT_RAW - EXACT_OUTPUT_TOTAL_RAW))"

if (( OUTPUT_ROUNDING_RAW < 0 )); then
  echo "Preview output is lower than requested exact output" >&2
  exit 1
fi

CWUSDC_BAL_RAW="$(cast call "${CWUSDC}" 'balanceOf(address)(uint256)' "${SIGNER}" --rpc-url "${ETHEREUM_MAINNET_RPC}" | awk '{print $1}')"
XAUT_BAL_RAW="$(cast call "${XAUT}" 'balanceOf(address)(uint256)' "${SIGNER}" --rpc-url "${ETHEREUM_MAINNET_RPC}" | awk '{print $1}')"
CWUSDC_ALLOWANCE_RAW="$(cast call "${CWUSDC}" 'allowance(address,address)(uint256)' "${SIGNER}" "${VAULT}" --rpc-url "${ETHEREUM_MAINNET_RPC}" | awk '{print $1}')"
XAUT_ALLOWANCE_RAW="$(cast call "${XAUT}" 'allowance(address,address)(uint256)' "${SIGNER}" "${VAULT}" --rpc-url "${ETHEREUM_MAINNET_RPC}" | awk '{print $1}')"
PROOF_USED="$(cast call "${VAULT}" 'usedProofIds(bytes32)(bool)' "${PROOF_ID}" --rpc-url "${ETHEREUM_MAINNET_RPC}" | tr -d '[:space:]')"

if [[ "${PROOF_USED}" == "true" ]]; then
  echo "Proof ID is already used: ${PROOF_ID}" >&2
  exit 1
fi
if (( CWUSDC_BAL_RAW < TOTAL_CWUSDC_IN_RAW )); then
  echo "Insufficient cWUSDC balance" >&2
  exit 1
fi
if (( XAUT_BAL_RAW < COLLATERAL_XAUT_RAW )); then
  echo "Insufficient XAUt balance" >&2
  exit 1
fi

CWUSDC_APPROVE_NEEDED=0
XAUT_APPROVE_NEEDED=0
if (( CWUSDC_ALLOWANCE_RAW < TOTAL_CWUSDC_IN_RAW )); then
  CWUSDC_APPROVE_NEEDED=1
fi
if (( XAUT_ALLOWANCE_RAW < COLLATERAL_XAUT_RAW )); then
  XAUT_APPROVE_NEEDED=1
fi

mkdir -p "$(dirname "${OUT_JSON}")"
python3 - "${OUT_JSON}" "${OUT_MD}" "${LATEST_JSON}" "${LATEST_MD}" \
  "${EXECUTE}" "${STAMP}" "${VAULT}" "${SIGNER}" "${RECIPIENT}" "${ROUNDING_RECEIVER}" \
  "${LOOPS}" "${EXACT_OUTPUT_PER_LOOP_RAW}" "${DEBT_RAW}" "${PROOF_ID}" "${ISO_HASH}" "${AUDIT_HASH}" "${PEG_HASH}" \
  "${COLLATERAL_XAUT_RAW}" "${CWUSDC_IN_PER_LOOP_RAW}" "${CWUSDC_OUT_PER_LOOP_RAW}" "${CWUSDC_LOSS_PER_LOOP_RAW}" \
  "${TOTAL_CWUSDC_IN_RAW}" "${TOTAL_CWUSDC_OUT_RAW}" "${EXACT_OUTPUT_TOTAL_RAW}" "${OUTPUT_ROUNDING_RAW}" "${TOTAL_NEUTRALIZED_RAW}" \
  "${CWUSDC_BAL_RAW}" "${XAUT_BAL_RAW}" "${CWUSDC_ALLOWANCE_RAW}" "${XAUT_ALLOWANCE_RAW}" "${CWUSDC_APPROVE_NEEDED}" "${XAUT_APPROVE_NEEDED}" <<'PY'
from decimal import Decimal
import json
from pathlib import Path
import sys

(
    out_json, out_md, latest_json, latest_md, execute, stamp, vault, signer, recipient, rounding_receiver,
    loops, exact_per_loop, debt, proof_id, iso_hash, audit_hash, peg_hash, collateral, cw_in, cw_out, loss,
    total_in, total_out, exact_total, rounding, neutralized, cw_bal, xaut_bal, cw_allow, xaut_allow,
    cw_approve, xaut_approve,
) = sys.argv[1:]

def units(raw):
    return f"{Decimal(int(raw)) / Decimal(10**6):f}"

payload = {
    "schema": "engine-x-loop-proof/v1",
    "executed": execute == "1",
    "stamp": stamp,
    "vault": vault,
    "signer": signer,
    "recipient": recipient,
    "roundingReceiver": rounding_receiver,
    "loops": int(loops),
    "debtUsdcPerLoopRaw": debt,
    "debtUsdcPerLoop": units(debt),
    "exactOutputPerLoopRaw": exact_per_loop,
    "exactOutputPerLoop": units(exact_per_loop),
    "proofId": proof_id,
    "isoHash": iso_hash,
    "auditHash": audit_hash,
    "pegHash": peg_hash,
    "preview": {
        "collateralXautRaw": collateral,
        "collateralXaut": units(collateral),
        "cwusdcInPerLoopRaw": cw_in,
        "cwusdcInPerLoop": units(cw_in),
        "cwusdcOutPerLoopRaw": cw_out,
        "cwusdcOutPerLoop": units(cw_out),
        "cwusdcLossPerLoopRaw": loss,
        "cwusdcLossPerLoop": units(loss),
        "totalCwusdcInRaw": total_in,
        "totalCwusdcIn": units(total_in),
        "totalCwusdcOutRaw": total_out,
        "totalCwusdcOut": units(total_out),
        "exactOutputTotalRaw": exact_total,
        "exactOutputTotal": units(exact_total),
        "outputRoundingRaw": rounding,
        "outputRounding": units(rounding),
        "totalNeutralizedRaw": neutralized,
        "totalNeutralized": units(neutralized),
    },
    "balances": {
        "cwusdcRaw": cw_bal,
        "cwusdc": units(cw_bal),
        "xautRaw": xaut_bal,
        "xaut": units(xaut_bal),
        "cwusdcAllowanceRaw": cw_allow,
        "xautAllowanceRaw": xaut_allow,
        "cwusdcApproveNeeded": cw_approve == "1",
        "xautApproveNeeded": xaut_approve == "1",
    },
}
Path(out_json).write_text(json.dumps(payload, indent=2) + "\n")
Path(latest_json).write_text(json.dumps(payload, indent=2) + "\n")
lines = [
    "# Engine X Loop Proof",
    "",
    f"- Executed: `{payload['executed']}`",
    f"- Vault: `{vault}`",
    f"- Loops: `{loops}`",
    f"- Exact cWUSDC output per loop: `{payload['exactOutputPerLoop']}`",
    f"- Debt USDC per loop: `{payload['debtUsdcPerLoop']}`",
    f"- cWUSDC in/out per loop: `{payload['preview']['cwusdcInPerLoop']} / {payload['preview']['cwusdcOutPerLoop']}`",
    f"- Neutralized per loop: `{payload['preview']['cwusdcLossPerLoop']}`",
    f"- Total exact output: `{payload['preview']['exactOutputTotal']}`",
    f"- Proof ID: `{proof_id}`",
    f"- cWUSDC approval needed: `{payload['balances']['cwusdcApproveNeeded']}`",
    f"- XAUt approval needed: `{payload['balances']['xautApproveNeeded']}`",
]
Path(out_md).write_text("\n".join(lines) + "\n")
Path(latest_md).write_text("\n".join(lines) + "\n")
PY

cat <<EOF
Engine X loop proof plan
  mode: ${EXECUTE}
  vault: ${VAULT}
  signer: ${SIGNER}
  recipient: ${RECIPIENT}
  loops: ${LOOPS}
  exact cWUSDC output per loop raw: ${EXACT_OUTPUT_PER_LOOP_RAW}
  debt USDC per loop raw: ${DEBT_RAW}
  cWUSDC in/out/loss per loop raw: ${CWUSDC_IN_PER_LOOP_RAW} / ${CWUSDC_OUT_PER_LOOP_RAW} / ${CWUSDC_LOSS_PER_LOOP_RAW}
  total exact output raw: ${EXACT_OUTPUT_TOTAL_RAW}
  total cWUSDC input raw: ${TOTAL_CWUSDC_IN_RAW}
  total neutralized raw: ${TOTAL_NEUTRALIZED_RAW}
  XAUt collateral raw: ${COLLATERAL_XAUT_RAW}
  proof id: ${PROOF_ID}
  cWUSDC approve needed: ${CWUSDC_APPROVE_NEEDED}
  XAUt approve needed: ${XAUT_APPROVE_NEEDED}
  report: ${OUT_JSON}
EOF

if [[ "${EXECUTE}" != "1" ]]; then
  cat <<EOF

Dry-run only. To broadcast this exact loop proof:
  EXECUTE=1 LOOPS=${LOOPS} EXACT_OUTPUT_PER_LOOP_RAW=${EXACT_OUTPUT_PER_LOOP_RAW} \\
    PROOF_ID=${PROOF_ID} ISO_HASH=${ISO_HASH} AUDIT_HASH=${AUDIT_HASH} PEG_HASH=${PEG_HASH} \\
    bash scripts/deployment/run-engine-x-loop-proof.sh
EOF
  exit 0
fi

if [[ "${CWUSDC_APPROVE_NEEDED}" == "1" ]]; then
  cast send "${CWUSDC}" 'approve(address,uint256)' "${VAULT}" "${TOTAL_CWUSDC_IN_RAW}" \
    --rpc-url "${ETHEREUM_MAINNET_RPC}" --private-key "${PRIVATE_KEY}"
fi
if [[ "${XAUT_APPROVE_NEEDED}" == "1" ]]; then
  cast send "${XAUT}" 'approve(address,uint256)' "${VAULT}" "${COLLATERAL_XAUT_RAW}" \
    --rpc-url "${ETHEREUM_MAINNET_RPC}" --private-key "${PRIVATE_KEY}"
fi

cast send "${VAULT}" \
  'runVirtualProofExactOutTo(bytes32,uint256,uint256,address,uint256,address,bytes32,bytes32,bytes32)' \
  "${PROOF_ID}" "${DEBT_RAW}" "${LOOPS}" "${RECIPIENT}" "${EXACT_OUTPUT_TOTAL_RAW}" "${ROUNDING_RECEIVER}" \
  "${ISO_HASH}" "${AUDIT_HASH}" "${PEG_HASH}" \
  --rpc-url "${ETHEREUM_MAINNET_RPC}" --private-key "${PRIVATE_KEY}"