import Foundation
import CryptoKit

/// Client (host) side of relay proof-of-possession enrollment (mesh P2). The relay room is
/// addressed by `roomId = sha256(host static public key)`, and that static key is *public* — it
/// rides in every pairing QR. So before the relay hands a host a credential for its room, the
/// host must prove it holds the *static private key* that derives the roomId, via an X25519
/// Diffie-Hellman challenge. This is the Swift counterpart of `cloud/nucleic-edge/src/relayEnroll.ts`;
/// the two must compute the same shared secret and proof byte-for-byte (pinned by a shared test
/// vector in `RelayEnrollmentTests`).
///
/// Flow:
///   1. host POSTs `{ staticKey: base64url(staticPub) }` → relay returns `{ challengeId, ephemeralPublicKey }`
///   2. host computes `shared = X25519(staticPriv, ephemeralPub)`,
///      `proof = HMAC-SHA256(shared, "nucleic-relay-enroll-v1:" + staticKeyB64)`
///   3. host POSTs `{ challengeId, proof }` → relay verifies and returns `{ hostId, hostSecret, roomId }`
public enum RelayEnrollment {
    /// The proof context string, byte-identical to the Worker's `PROOF_CONTEXT`.
    static let proofContext = "nucleic-relay-enroll-v1:"

    /// The exact `staticKey` string the host sends in the challenge request and the Worker folds
    /// into the proof — base64url of the static public key, no padding.
    public static func staticKeyEncoding(_ staticPublicKey: Data) -> String {
        base64URLNoPad(staticPublicKey)
    }

    /// Compute the enrollment proof for a received ephemeral public key. `staticKeyB64` must be the
    /// exact string sent in the challenge (`staticKeyEncoding(identity.staticPublicKey)`), because
    /// the Worker binds the proof to the string it received verbatim.
    public static func proof(
        staticKey: Curve25519.KeyAgreement.PrivateKey,
        ephemeralPublicKey: Data,
        staticKeyB64: String
    ) throws -> Data {
        let ephPub = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: ephemeralPublicKey)
        let shared = try staticKey.sharedSecretFromKeyAgreement(with: ephPub)
        let sharedBytes = shared.withUnsafeBytes { Data($0) }
        let message = Data((proofContext + staticKeyB64).utf8)
        let mac = HMAC<SHA256>.authenticationCode(for: message, using: SymmetricKey(data: sharedBytes))
        return Data(mac)
    }
}

/// base64url without padding, matching Node's `Buffer.toString("base64url")` and the Worker's
/// `b64urlEncode` — so the `staticKey` string is identical on both sides of the proof.
private func base64URLNoPad(_ data: Data) -> String {
    data.base64EncodedString()
        .replacingOccurrences(of: "+", with: "-")
        .replacingOccurrences(of: "/", with: "_")
        .replacingOccurrences(of: "=", with: "")
}
