import { BigNumber, utils } from 'ethers'

import BalanceTree from './balance-tree'

const { isAddress, getAddress } = utils

// This is the blob that gets distributed and pinned to IPFS.
// It is completely sufficient for recreating the entire merkle tree.
// Anyone can verify that all air drops are included in the tree,
// and the tree has no additional distributions.
export interface MerkleDistributorInfo {
  merkleRoot: string
  tokenTotal: string
  claims: {
    [account: string]: {
      index: number
      amount: string
      proof: string[]
      flags?: {
        [flag: string]: boolean
      }
    }
  }
}

type OldFormat = { [account: string]: number | string }
type NewFormat = { address: string; earnings: string; reasons: string }

export const parseBalanceMap = (
  balances: OldFormat | NewFormat[]
): MerkleDistributorInfo => {
  // if balances are in an old format, process them
  const balancesInNewFormat: NewFormat[] = Array.isArray(balances)
    ? balances
    : Object.keys(balances).map((account): NewFormat => {
        let earnings: string
        if (typeof balances[account] === 'number') {
          earnings = `0x${balances[account].toString(16)}`
        } else {
          earnings = BigNumber.from(balances[account]).toHexString()
        }

        return {
          address: account,
          earnings,
          reasons: '',
        }
      })

  const dataByAddress = balancesInNewFormat.reduce<{
    [address: string]: {
      amount: BigNumber
      flags?: { [flag: string]: boolean }
    }
  }>((memo, { address: account, earnings, reasons }) => {
    if (!isAddress(account)) {
      throw new Error(`Found invalid address: ${account}`)
    }
    const parsed = getAddress(account)
    if (memo[parsed]) {
      throw new Error(`Duplicate address: ${parsed}`)
    }
    const parsedNum = BigNumber.from(earnings)
    if (parsedNum.lte(0)) {
      throw new Error(`Invalid amount for account: ${account}`)
    }

    const flags = {
      isSOCKS: reasons.includes('socks'),
      isLP: reasons.includes('lp'),
      isUser: reasons.includes('user'),
    }

    memo[parsed] = { amount: parsedNum, ...(reasons === '' ? {} : { flags }) }
    return memo
  }, {})

  const sortedAddresses = Object.keys(dataByAddress).sort()

  // construct a tree
  const tree = new BalanceTree(
    sortedAddresses.map((address) => ({
      account: address,
      amount: dataByAddress[address].amount,
    }))
  )

  // generate claims
  const claims = sortedAddresses.reduce<{
    [address: string]: {
      amount: string
      index: number
      proof: string[]
      flags?: { [flag: string]: boolean }
    }
  }>((memo, address, index) => {
    const { amount, flags } = dataByAddress[address]
    memo[address] = {
      index,
      amount: amount.toHexString(),
      proof: tree.getProof(index, address, amount),
      ...(flags ? { flags } : {}),
    }
    return memo
  }, {})

  const tokenTotal: BigNumber = sortedAddresses.reduce<BigNumber>(
    (memo, key) => memo.add(dataByAddress[key].amount),
    BigNumber.from(0)
  )

  return {
    merkleRoot: tree.getHexRoot(),
    tokenTotal: tokenTotal.toHexString(),
    claims,
  }
}
