import Foundation

#if canImport(CryptoKit)
import CryptoKit
#else
import Crypto
#endif

/// The relay routing envelope (mesh P2), the Swift counterpart of the Room DO's framing in
/// `cloud/nucleic-edge/src/room.ts`. It exists **only on the host ↔ relay leg**: the relay stamps
/// a sender tag onto client→host frames so the host can demux which peer sent them, and the host
/// prepends a target tag onto host→client frames so the relay can route them to exactly one
/// client. The client ↔ relay leg is always raw Noise frames, so Noise and everything above see
/// identical bytes on every transport.
///
/// Wire shape: `[0x01 version][8-byte tag = sha256(deviceID)[0..<8]][raw Noise frame]`. The host
/// never needs a peer's plaintext deviceID for routing — the tag round-trips (the relay maps
/// tag → client socket), so `RelayListener` keys its per-peer virtual channels by the raw tag.
public enum RelayEnvelope {
    public static let version: UInt8 = 0x01
    public static let tagLength = 8
    public static let headerLength = 1 + 8 // version + tag

    /// The routing tag for a device id: the first 8 bytes of its SHA-256. Matches `deviceTag()`
    /// in room.ts (`sha256(deviceId).slice(0, 8)`).
    public static func tag(forDeviceID deviceID: String) -> Data {
        Data(SHA256.hash(data: Data(deviceID.utf8)).prefix(tagLength))
    }

    /// Wrap a raw frame for the addressed tag (host → relay).
    public static func wrap(_ frame: Data, tag: Data) -> Data {
        precondition(tag.count == tagLength, "relay tag must be \(tagLength) bytes")
        var out = Data(capacity: headerLength + frame.count)
        out.append(version)
        out.append(tag)
        out.append(frame)
        return out
    }

    /// Split an enveloped frame into its tag + raw payload (relay → host). Returns nil for a
    /// malformed or wrong-version envelope, mirroring the DO's drop-on-malformed behavior.
    public static func unwrap(_ enveloped: Data) -> (tag: Data, frame: Data)? {
        guard enveloped.count >= headerLength, enveloped.first == version else { return nil }
        let base = enveloped.startIndex
        let tag = enveloped.subdata(in: base.advanced(by: 1)..<base.advanced(by: headerLength))
        let frame = enveloped.subdata(in: base.advanced(by: headerLength)..<enveloped.endIndex)
        return (tag, frame)
    }
}
