import { address } from "bitcoinjs-lib"
import {
  ByteArray,
  concat,
  encodePacked,
  Hex,
  ripemd160,
  sha256,
  toBytes,
  toHex,
} from "viem"
import bs58 from "bs58"

const supportedAddresses = {
  P2WPKH: {
    prefixes: ["bc1", "tb1"],
  },
  P2PKH: {
    prefixes: ["1", "m", "n"],
  },
  P2SH_P2WPKH: {
    prefixes: ["3", "2"],
  },
}

const networkAddressPrefixes = {
  mainnet: ["bc1", "1", "3"],
  testnet: ["tb1", "m", "n", "2"],
}

export class UnsupportedBitcoinAddressError extends Error {
  constructor() {
    const message = `Unsupported address. Supported addresses: ${Object.keys(supportedAddresses).join(", ")}`
    super(message)
  }
}

function checkIfAddressStartsWith(bitcoinAddress: string, prefixes: string[]) {
  return prefixes.some((prefix) =>
    bitcoinAddress.toLowerCase().startsWith(prefix),
  )
}

// HASH160(publicKey))
function recoverPubkeyHashP2PKH(bitcoinBase58Address: string): Hex {
  const { hash, version } = address.fromBase58Check(bitcoinBase58Address)

  // From docs
  // https://bitcoinjs.github.io/bitcoinjs-lib/interfaces/address.Base58CheckResult.html#version
  // Address version: 0x00 for P2PKH, 0x05 for P2SH
  // Testnet address version: 111 for P2PKH, 196 for P2SH
  if (version === 0 || version === 111) return toHex(hash)

  throw new UnsupportedBitcoinAddressError()
}

// HASH160(0014 | HASH160(publicKey))
function recoverRedeemScriptHashNestedSegwit(
  bitcoinBase58Address: string,
): Hex {
  const { hash, version } = address.fromBase58Check(bitcoinBase58Address)

  if (version === 5 || version === 196) {
    return toHex(hash)
  }

  throw new UnsupportedBitcoinAddressError()
}

// HASH160(publicKey))
function recoverPubKeyHashP2WPKH(bitcoinBech32Address: string): Hex {
  const { data } = address.fromBech32(bitcoinBech32Address)

  // From docs
  // https://bitcoinjs.github.io/bitcoinjs-lib/interfaces/address.Bech32Result.html#data:
  // Address data：20 bytes for P2WPKH; 32 bytes for P2WSH, P2TR
  if (data.length !== 20) throw new UnsupportedBitcoinAddressError()

  return toHex(data)
}

// Takes the compressed public key and outputs P2SH-P2WPKH address for that
// public key as specified in BIP-141.
function toNestedSegwitAddress(
  compressedPublicKey: string | ByteArray,
  isTestnet = false,
): string {
  if (compressedPublicKey.length !== 66) {
    // According to BIP-143 only compressed public keys are accepted for SegWit
    // addresses. Sending to address generated from uncompressed public key will
    // result in losing tokens. The compressed public key should have 66 bytes:
    // 32 bytes for X and 1 byte for Y.
    throw new Error(
      `The value [${compressedPublicKey}] is not a compressed public key`,
    )
  }

  const pubKeyHash = ripemd160(sha256(`0x${compressedPublicKey}`))
  const scriptPubKey = ripemd160(
    sha256(encodePacked(["bytes2", "bytes20"], ["0x0014", pubKeyHash])),
  )

  const prefix = isTestnet ? "0xc4" : "0x05"

  const versionedData = concat([prefix, scriptPubKey])
  // Checksum is the first 4 bytes of the double-sha256 of the versioned data.
  const checksum = toBytes(sha256(sha256(versionedData))).slice(0, 4)

  // The parameters passed to the `concat` function must be of the same type.
  // Therefore, we need to convert `versionedData` into bytes form.
  return bs58.encode(concat([toBytes(versionedData), checksum]))
}

/**
 * Checks if the Bitcoin address is really nested segwit (`P2SH-P2WPKH`).
 * @param bitcoinAddress Bitcoin address to check.
 * @param publicKey Public key of the bitcoin address.
 * @returns True if the address is nested segwit address (`P2SH-P2WPKH`),
 *          otherwise false.
 */
function isNestedSegwitAddress(bitcoinAddress: string, publicKey?: string) {
  const hasCorrectPrefix = checkIfAddressStartsWith(
    bitcoinAddress,
    supportedAddresses.P2SH_P2WPKH.prefixes,
  )

  const isTestnet = bitcoinAddress.startsWith("2")

  try {
    return (
      hasCorrectPrefix &&
      publicKey &&
      toNestedSegwitAddress(
        publicKey?.startsWith("0x") ? publicKey.slice(2) : publicKey,
        isTestnet,
      ) === bitcoinAddress
    )
  } catch (error) {
    return false
  }
}

// Important: If the address is P2SH, the function requires the public key.
// There is no way to validate if P2SH address is P2SH-P2WPKH without knowing
// the public key hash. It is recommended to use `toNestedSegwitAddress`
// function first to validate if the address is really P2SH-P2WPKH.
function recoverTruncatedBitcoinAddress(
  bitcoinAddress: string,
  publicKey?: string,
): Hex {
  const isP2PKH = checkIfAddressStartsWith(
    bitcoinAddress,
    supportedAddresses.P2PKH.prefixes,
  )

  if (isP2PKH) {
    return recoverPubkeyHashP2PKH(bitcoinAddress)
  }

  const isP2WPKH = checkIfAddressStartsWith(
    bitcoinAddress,
    supportedAddresses.P2WPKH.prefixes,
  )

  if (isP2WPKH) {
    return recoverPubKeyHashP2WPKH(bitcoinAddress)
  }

  if (isNestedSegwitAddress(bitcoinAddress, publicKey)) {
    return recoverRedeemScriptHashNestedSegwit(bitcoinAddress)
  }

  throw new UnsupportedBitcoinAddressError()
}

function isP2WPKHAddress(bitcoinAddress: string) {
  try {
    recoverPubKeyHashP2WPKH(bitcoinAddress)
    return true
  } catch (err) {
    return false
  }
}

function isP2PKHAddress(bitcoinAddress: string) {
  try {
    recoverPubkeyHashP2PKH(bitcoinAddress)
    return true
  } catch (err) {
    return false
  }
}

function isMainnetAddress(bitcoinAddress: string) {
  return checkIfAddressStartsWith(
    bitcoinAddress,
    networkAddressPrefixes.mainnet,
  )
}

function isTestnetAddress(bitcoinAddress: string) {
  return checkIfAddressStartsWith(
    bitcoinAddress,
    networkAddressPrefixes.testnet,
  )
}

export default {
  toNestedSegwitAddress,
  recoverTruncatedBitcoinAddress,
  isP2WPKHAddress,
  isP2PKHAddress,
  isNestedSegwitAddress,
  isMainnetAddress,
  isTestnetAddress,
}