import Foundation
import Security
import CryptoKit

public enum SecureEnclaveKeyManagerError: Error, LocalizedError {
  case generateFailed(String)
  case exportFailed(String)
  case keyNotFound(String)
  case signFailed(String)
  case invalidBase64Input

  public var errorDescription: String? {
    switch self {
    case .generateFailed(let detail):
      return "Failed to generate key pair: \(detail)"
    case .exportFailed(let detail):
      return "Failed to export public key: \(detail)"
    case .keyNotFound(let tag):
      return "Key not found for tag: \(tag)"
    case .signFailed(let detail):
      return "Failed to sign: \(detail)"
    case .invalidBase64Input:
      return "Invalid base64url-encoded input"
    }
  }
}

/// Platform-agnostic Secure Enclave key manager.
/// Provides P-256 key generation, signing, and management backed by the iOS Secure Enclave.
/// All public keys are returned in uncompressed SEC1 format (65 bytes: 04 || x || y).
/// All binary data uses base64url encoding (RFC 4648 §5, no padding).
public class SecureEnclaveKeyManager {

  // No state needed — all operations use the system Keychain and Secure Enclave directly
  public init() {}

  public func isAvailable() -> Bool {
    if #available(iOS 13.0, *) {
      return SecureEnclave.isAvailable
    }
    return false
  }

  public func hasKey(tag: String) -> Bool {
    let query = baseQuery(for: tag, returning: true)
    var item: CFTypeRef?
    let status = SecItemCopyMatching(query as CFDictionary, &item)
    return status == errSecSuccess
  }

  public func generateKeyPair(tag: String) throws -> String {
    let tagData = tag.data(using: .utf8)!

    let access = SecAccessControlCreateWithFlags(
      kCFAllocatorDefault,
      kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly,
      .privateKeyUsage,
      nil
    )!

    let attributes: [String: Any] = [
      kSecAttrKeyType as String: kSecAttrKeyTypeECSECPrimeRandom,
      kSecAttrKeySizeInBits as String: 256,
      kSecAttrTokenID as String: kSecAttrTokenIDSecureEnclave,
      kSecPrivateKeyAttrs as String: [
        kSecAttrIsPermanent as String: true,
        kSecAttrApplicationTag as String: tagData,
        kSecAttrAccessControl as String: access,
      ] as [String: Any],
    ]

    var error: Unmanaged<CFError>?
    guard let privateKey = SecKeyCreateRandomKey(attributes as CFDictionary, &error) else {
      throw SecureEnclaveKeyManagerError.generateFailed(
        "\(error!.takeRetainedValue())"
      )
    }

    guard let publicKey = SecKeyCopyPublicKey(privateKey) else {
      throw SecureEnclaveKeyManagerError.exportFailed("Failed to extract public key")
    }
    let publicKeyData = try exportPublicKey(publicKey)
    return base64urlEncode(publicKeyData)
  }

  public func getPublicKey(tag: String) throws -> String? {
    let query = baseQuery(for: tag, returning: true)

    var item: CFTypeRef?
    let status = SecItemCopyMatching(query as CFDictionary, &item)

    guard status == errSecSuccess else {
      return nil
    }

    guard let item else {
      throw SecureEnclaveKeyManagerError.keyNotFound(tag)
    }
    let privateKey = item as! SecKey
    guard let publicKey = SecKeyCopyPublicKey(privateKey) else {
      throw SecureEnclaveKeyManagerError.exportFailed("Failed to extract public key")
    }
    let publicKeyData = try exportPublicKey(publicKey)
    return base64urlEncode(publicKeyData)
  }

  public func sign(tag: String, payload: Data) throws -> String {
    let query = baseQuery(for: tag, returning: true)

    var item: CFTypeRef?
    let status = SecItemCopyMatching(query as CFDictionary, &item)

    guard status == errSecSuccess else {
      throw SecureEnclaveKeyManagerError.keyNotFound(tag)
    }

    guard let item else {
      throw SecureEnclaveKeyManagerError.keyNotFound(tag)
    }
    let privateKey = item as! SecKey

    var error: Unmanaged<CFError>?
    guard let signature = SecKeyCreateSignature(
      privateKey,
      .ecdsaSignatureMessageX962SHA256,
      payload as CFData,
      &error
    ) else {
      throw SecureEnclaveKeyManagerError.signFailed(
        "\(error!.takeRetainedValue())"
      )
    }

    return base64urlEncode(signature as Data)
  }

  public func deleteKey(tag: String) {
    let query = baseQuery(for: tag, returning: false)
    SecItemDelete(query as CFDictionary)
  }

  // MARK: - Private helpers

  private func baseQuery(for tag: String, returning returnRef: Bool) -> [String: Any] {
    var query: [String: Any] = [
      kSecClass as String: kSecClassKey,
      kSecAttrApplicationTag as String: tag.data(using: .utf8)!,
      kSecAttrKeyType as String: kSecAttrKeyTypeECSECPrimeRandom,
    ]
    if returnRef {
      query[kSecReturnRef as String] = true
    }
    return query
  }

  private func exportPublicKey(_ publicKey: SecKey) throws -> Data {
    var error: Unmanaged<CFError>?
    guard let data = SecKeyCopyExternalRepresentation(publicKey, &error) else {
      throw SecureEnclaveKeyManagerError.exportFailed(
        "\(error!.takeRetainedValue())"
      )
    }
    // SecKeyCopyExternalRepresentation returns uncompressed SEC1 format (04 || x || y)
    return data as Data
  }

  private func base64urlEncode(_ data: Data) -> String {
    data.base64EncodedString()
      .replacingOccurrences(of: "+", with: "-")
      .replacingOccurrences(of: "/", with: "_")
      .replacingOccurrences(of: "=", with: "")
  }
}

public func base64urlDecode(_ string: String) throws -> Data {
  var base64 = string
    .replacingOccurrences(of: "-", with: "+")
    .replacingOccurrences(of: "_", with: "/")
  let remainder = base64.count % 4
  if remainder > 0 {
    base64 += String(repeating: "=", count: 4 - remainder)
  }
  guard let data = Data(base64Encoded: base64) else {
    throw SecureEnclaveKeyManagerError.invalidBase64Input
  }
  return data
}