import CryptoKit
import Foundation

/// BIP-32 key derivation for the XRP account path `m/44'/144'/0'/0/0`. The
/// first three steps are hardened, the last two are not, so both cases are
/// supported (hardened steps sign with the parent private key, normal steps
/// with the parent public key).
enum BIP32 {
  static let hardenedBit: UInt32 = 0x8000_0000

  /// Standard XRP derivation path: 44' / 144' / 0' / 0 / 0.
  static let xrpPath: [UInt32] = [
    hardenedBit | 44,
    hardenedBit | 144,
    hardenedBit | 0,
    0,
    0,
  ]

  struct Key {
    var privateKey: [UInt8]
    var chainCode: [UInt8]
    var publicKey: [UInt8]
  }

  static func masterKey(fromSeed seed: [UInt8]) -> Key? {
    let derived = hmac(key: Array("Bitcoin seed".utf8), data: seed)
    guard derived.count == 64 else { return nil }
    let privateKey = Array(derived[0..<32])
    guard isValidScalar(privateKey),
          let publicKey = Secp256k1Backend.publicKey?(Data(privateKey)) else { return nil }
    return Key(privateKey: privateKey,
               chainCode: Array(derived[32..<64]),
               publicKey: Array(publicKey))
  }

  static func derive(_ path: [UInt32], from key: Key) -> Key? {
    var current = key
    for index in path {
      guard let next = deriveChild(current, index: index) else { return nil }
      current = next
    }
    return current
  }

  private static func deriveChild(_ parent: Key, index: UInt32) -> Key? {
    let isHardened = index & hardenedBit != 0

    var data: [UInt8] = []
    if isHardened {
      data.append(0x00)
      data += parent.privateKey
    } else {
      data += parent.publicKey
    }
    data += [
      UInt8((index >> 24) & 0xff),
      UInt8((index >> 16) & 0xff),
      UInt8((index >> 8) & 0xff),
      UInt8(index & 0xff),
    ]

    let derived = hmac(key: parent.chainCode, data: data)
    guard derived.count == 64 else { return nil }
    let left = Array(derived[0..<32])
    guard isValidScalar(left),
          let sum = Secp256k1Backend.addScalar?(Data(left), Data(parent.privateKey)) else { return nil }
    let privateKey = Array(sum)
    // The child public key is the same whether it is computed from the child
    // private key or by point addition, so derive it directly.
    guard isValidScalar(privateKey),
          let publicKey = Secp256k1Backend.publicKey?(Data(privateKey)) else { return nil }

    return Key(privateKey: privateKey,
               chainCode: Array(derived[32..<64]),
               publicKey: Array(publicKey))
  }

  /// A key is usable when 0 < key < the secp256k1 group order.
  private static func isValidScalar(_ bytes: [UInt8]) -> Bool {
    guard bytes.count == 32 else { return false }
    let orderHex = "fffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141"
    guard let order = Hex.decode(orderHex) else { return false }
    if bytes.allSatisfy({ $0 == 0 }) { return false }
    for (left, right) in zip(bytes, order) where left != right { return left < right }
    return false
  }

  private static func hmac(key: [UInt8], data: [UInt8]) -> [UInt8] {
    Array(HMAC<SHA512>.authenticationCode(for: Data(data), using: SymmetricKey(data: Data(key))))
  }
}