import CryptoKit
import Foundation

/// BIP-39 mnemonics. The XRP wallet uses a 12 word phrase as its recovery
/// secret, matching what hardware and mobile wallets expect.
enum BIP39 {
  static let wordCount = 12
  static let entropyBytes = 16
  private static let checksumBits = entropyBytes / 4

  private static let index: [String: Int] = {
    var map: [String: Int] = [:]
    for (position, word) in BIP39Wordlist.words.enumerated() { map[word] = position }
    return map
  }()

  static func generateMnemonic() -> String {
    var entropy = [UInt8](repeating: 0, count: entropyBytes)
    for position in entropy.indices { entropy[position] = UInt8.random(in: 0...255) }
    return mnemonic(from: entropy)
  }

  /// 16 bytes of entropy plus its 4 bit SHA-256 checksum, as 12 words.
  static func mnemonic(from entropy: [UInt8]) -> String {
    precondition(entropy.count == entropyBytes, "BIP-39 entropy must be 16 bytes")
    var bits = entropy.flatMap { byte in (0..<8).map { (byte >> (7 - $0)) & 1 } }
    let checksum = XRPHash.sha256(entropy)
    bits += (0..<checksumBits).map { (checksum[0] >> (7 - $0)) & 1 }

    var words: [String] = []
    for start in stride(from: 0, to: bits.count, by: 11) {
      var value = 0
      for offset in 0..<11 { value = (value << 1) | Int(bits[start + offset]) }
      words.append(BIP39Wordlist.words[value])
    }
    return words.joined(separator: " ")
  }

  static func entropy(from mnemonic: String) -> [UInt8]? {
    let words = normalize(mnemonic).split(separator: " ").map(String.init)
    guard words.count == wordCount else { return nil }

    var bits: [Int] = []
    for word in words {
      guard let position = index[word] else { return nil }
      bits += (0..<11).map { (position >> (10 - $0)) & 1 }
    }

    var entropy: [UInt8] = []
    for start in stride(from: 0, to: entropyBytes * 8, by: 8) {
      var value: UInt8 = 0
      for offset in 0..<8 { value = (value << 1) | UInt8(bits[start + offset]) }
      entropy.append(value)
    }

    let checksum = XRPHash.sha256(entropy)
    for offset in 0..<checksumBits {
      let expected = (checksum[0] >> (7 - offset)) & 1
      guard bits[entropyBytes * 8 + offset] == expected else { return nil }
    }
    return entropy
  }

  /// 64 byte BIP-39 seed (PBKDF2-HMAC-SHA512, empty passphrase).
  static func seed(from mnemonic: String) -> [UInt8] {
    let password = Array(normalize(mnemonic).utf8)
    let salt = Array("mnemonic".utf8)
    // Derived length equals the hash length, so a single PBKDF2 block is needed.
    let block = salt + [0, 0, 0, 1]
    var previous = hmacSHA512(key: password, data: block)
    var result = previous
    for _ in 1..<2_048 {
      previous = hmacSHA512(key: password, data: previous)
      for position in result.indices { result[position] ^= previous[position] }
    }
    return result
  }

  static func isMnemonic(_ secret: String) -> Bool {
    normalize(secret).split(separator: " ").count == wordCount
  }

  private static func normalize(_ text: String) -> String {
    text.trimmingCharacters(in: .whitespacesAndNewlines)
      .lowercased()
      .precomposedStringWithCompatibilityMapping
  }

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