const { ethers } = require('hardhat');
const { getStorageAt, setStorageAt } = require('@nomicfoundation/hardhat-network-helpers');

const ImplementationLabel = 'eip1967.proxy.implementation';
const AdminLabel = 'eip1967.proxy.admin';
const BeaconLabel = 'eip1967.proxy.beacon';

const erc1967slot = label => ethers.toBeHex(ethers.toBigInt(ethers.id(label)) - 1n);
const erc7201slot = label => ethers.toBeHex(ethers.toBigInt(ethers.keccak256(erc1967slot(label))) & ~0xffn);

const getSlot = (address, slot) =>
  (ethers.isAddressable(address) ? address.getAddress() : Promise.resolve(address)).then(address =>
    getStorageAt(address, ethers.isBytesLike(slot) ? slot : erc1967slot(slot)),
  );

const setSlot = (address, slot, value) =>
  Promise.all([
    ethers.isAddressable(address) ? address.getAddress() : Promise.resolve(address),
    ethers.isAddressable(value) ? value.getAddress() : Promise.resolve(value),
  ]).then(([address, value]) => setStorageAt(address, ethers.isBytesLike(slot) ? slot : erc1967slot(slot), value));

const getAddressInSlot = (address, slot) =>
  getSlot(address, slot).then(slotValue => ethers.AbiCoder.defaultAbiCoder().decode(['address'], slotValue)[0]);

module.exports = {
  ImplementationLabel,
  AdminLabel,
  BeaconLabel,
  ImplementationSlot: erc1967slot(ImplementationLabel),
  AdminSlot: erc1967slot(AdminLabel),
  BeaconSlot: erc1967slot(BeaconLabel),
  erc1967slot,
  erc7201slot,
  setSlot,
  getSlot,
  getAddressInSlot,
};