import { PublicActions } from 'viem';
import { Address } from '../types';
import { getERC4337Nonce } from './erc4337nonce';
import { getPermit2Nonce } from './permit2-nonce';
import { getPermitNonce } from './permit-nonce';

const REPLAYABLE_NONCE_KEY = 8453n;

export class NonceManager {
  private erc4337Nonce: bigint = 0n;
  protected permit2Nonce: bigint = 0n;
  protected permitNonceMap: Record<Address, bigint> = {};

  protected publicActions: PublicActions;
  private accountAddress?: Address;

  constructor(publicActions: PublicActions) {
    this.erc4337Nonce = 0n;
    this.permit2Nonce = 0n;
    this.publicActions = publicActions;
  }

  setAccountAddress(accountAddress: Address) {
    this.accountAddress = accountAddress;
  }

  public async getNonce(
    from?: Address,
    withoutChainIdValidation: boolean = false
  ) {
    const sender = from || this.accountAddress;

    if (!sender) {
      throw new Error('account not initialized');
    }

    const nonce = await getERC4337Nonce(
      this.publicActions,
      sender,
      withoutChainIdValidation ? REPLAYABLE_NONCE_KEY : BigInt(0)
    );

    this.erc4337Nonce = nonce;

    return nonce;
  }

  public async updatePermit2Nonce(spender: Address, from?: Address) {
    const sender = from || this.accountAddress;

    if (!sender) {
      throw new Error('account not initialized');
    }

    const nonce = await getPermit2Nonce(this.publicActions, sender, spender);

    this.permit2Nonce = nonce;

    return nonce;
  }

  public async updatePermitNonce(tokenAddress: Address) {
    if (!this.accountAddress) {
      throw new Error('account not initialized');
    }

    const nonce = await getPermitNonce(
      this.publicActions,
      tokenAddress,
      this.accountAddress
    );

    this.permitNonceMap[tokenAddress] = nonce;

    return nonce;
  }

  getPermitNonce(tokenAddress: Address) {
    return this.permitNonceMap[tokenAddress];
  }

  getPermit2Nonce() {
    return this.permit2Nonce;
  }

  getERC4337Nonce() {
    return this.erc4337Nonce;
  }
}