import CryptoKit
import Foundation

/// The relay envelope: what is signed, what is encrypted, and in what order.
///
/// This is the third implementation of a format that Go and Node already implement, and the three
/// share no code on purpose. `contracts/test-vectors.json` is recomputed here by `RelayWireTests`,
/// so a byte-level disagreement shows up as a failing test rather than as a signature failure in
/// the field — where the only symptom is "the node refused the request" with no way to tell which
/// of the three encoders drifted.
///
/// The relay is not trusted, so signing and encryption are separate concerns: Ed25519 covers a
/// canonical byte encoding of the routing fields plus the ciphertext, which stops a compromised
/// relay from retargeting a request at another node or swapping a response between two in-flight
/// requests. AES-256-GCM then carries the payload under keys from an ephemeral X25519 exchange.
public enum RelayWire {
    public static let version: UInt64 = 1
    public static let suite = "X25519-ED25519-AES256GCM-HKDFSHA256"

    static let domainRequest = "msa-relay-v1-req"
    static let domainResponse = "msa-relay-v1-resp"
    private static let label = "msa-relay-v1"

    public static let keyBytes = 32
    public static let ivBytes = 12
    public static let tagBytes = 16

    public enum Failure: Error, Equatable {
        case unsupportedVersion(UInt64)
        case unsupportedSuite(String)
        case badKeyLength(field: String, got: Int)
        case missingIdentityField
        case signatureDoesNotVerify
        case ciphertextFailedToAuthenticate
        case peerKeyUnusable
        case invalidResponseBody
    }

    // MARK: - Envelope

    /// The outer, relay-visible structure. The relay routes on `node` and correlates on `req`, and
    /// can read nothing else: `ek` is a public key, `p` is ciphertext, and `s` covers everything.
    public struct Envelope: Codable, Equatable, Sendable {
        public var v: UInt64
        public var node: String
        public var req: String
        /// The paired device id. Present on requests only: it is signed, so the relay cannot
        /// relabel one phone's traffic as another's, and the agent uses it to pick the key to
        /// verify against.
        public var app: String?
        public var ts: UInt64
        /// X25519 public key: the phone's per-request ephemeral on a request, the agent's
        /// per-response ephemeral on a response.
        public var ek: Data
        public var alg: String
        /// AES-256-GCM ciphertext with the 16-byte tag appended.
        public var p: Data
        public var s: Data?

        public init(
            v: UInt64 = RelayWire.version,
            node: String,
            req: String,
            app: String? = nil,
            ts: UInt64,
            ek: Data,
            alg: String = RelayWire.suite,
            p: Data,
            s: Data? = nil
        ) {
            self.v = v
            self.node = node
            self.req = req
            self.app = app
            self.ts = ts
            self.ek = ek
            self.alg = alg
            self.p = p
            self.s = s
        }

        /// The wire form is base64 for the byte fields, and the key set is fixed: the relay parses
        /// these same names, and the Go agent's struct tags are the authority for them.
        private enum CodingKeys: String, CodingKey {
            case v, node, req, app, ts, ek, alg, p, s
        }

        public init(from decoder: Decoder) throws {
            let container = try decoder.container(keyedBy: CodingKeys.self)
            v = try container.decode(UInt64.self, forKey: .v)
            node = try container.decode(String.self, forKey: .node)
            req = try container.decode(String.self, forKey: .req)
            app = try container.decodeIfPresent(String.self, forKey: .app)
            ts = try container.decode(UInt64.self, forKey: .ts)
            alg = try container.decode(String.self, forKey: .alg)
            ek = try Self.decodeBase64(container.decode(String.self, forKey: .ek), field: "ek")
            p = try Self.decodeBase64(container.decode(String.self, forKey: .p), field: "p")
            let signature = try container.decodeIfPresent(String.self, forKey: .s) ?? ""
            s = signature.isEmpty ? nil : try Self.decodeBase64(signature, field: "s")
        }

        public func encode(to encoder: Encoder) throws {
            var container = encoder.container(keyedBy: CodingKeys.self)
            try container.encode(v, forKey: .v)
            try container.encode(node, forKey: .node)
            try container.encode(req, forKey: .req)
            try container.encodeIfPresent(app, forKey: .app)
            try container.encode(ts, forKey: .ts)
            try container.encode(ek.base64EncodedString(), forKey: .ek)
            try container.encode(alg, forKey: .alg)
            try container.encode(p.base64EncodedString(), forKey: .p)
            if let s {
                try container.encode(s.base64EncodedString(), forKey: .s)
            }
        }

        /// Strict base64: Go's decoder silently accepts some non-canonical forms, and a signature
        /// must be computed over exactly the bytes the sender used, so a lenient decode here would
        /// turn a malformed envelope into a valid-looking one.
        private static func decodeBase64(_ value: String, field: String) throws -> Data {
            guard let data = Data(base64Encoded: value, options: []) else {
                throw Failure.invalidResponseBody
            }
            return data
        }
    }

    /// The decrypted request payload.
    public struct InnerRequest: Codable, Equatable, Sendable {
        /// Upstream method. The agent matches it against its policy before dialing anything.
        public var m: String
        /// Upstream path, already carrying any query string. The agent forwards it verbatim and
        /// never assembles one from parts, so there is no path-traversal surface from the relay.
        public var p: String
        /// sha256 over the method, path and canonical body, produced by an explicit user gesture.
        /// Mandatory for every entry in the agent's `dangerous` list and absent for reads.
        public var intent: Data?
        /// Raw upstream body, forwarded unchanged. Kept as bytes so the agent never re-serialises
        /// JSON and thereby never changes what the intent was computed over.
        public var body: Data?

        public init(m: String, p: String, intent: Data? = nil, body: Data? = nil) {
            self.m = m
            self.p = p
            self.intent = intent
            self.body = body
        }

        private enum CodingKeys: String, CodingKey { case m, p, intent, body }

        public init(from decoder: Decoder) throws {
            let container = try decoder.container(keyedBy: CodingKeys.self)
            m = try container.decode(String.self, forKey: .m)
            p = try container.decode(String.self, forKey: .p)
            intent = try container.decodeIfPresent(Data.self, forKey: .intent)
            body = try container.decodeIfPresent(Data.self, forKey: .body)
        }

        public func encode(to encoder: Encoder) throws {
            var container = encoder.container(keyedBy: CodingKeys.self)
            try container.encode(m, forKey: .m)
            try container.encode(p, forKey: .p)
            if let intent { try container.encode(intent, forKey: .intent) }
            if let body { try container.encode(body, forKey: .body) }
        }
    }

    /// The decrypted response payload. Status lives inside the ciphertext so the relay cannot
    /// rewrite a failure into a success.
    public struct InnerResponse: Codable, Equatable, Sendable {
        public var status: Int
        /// The agent's own classification of what happened. Inside the ciphertext for the same
        /// reason as `status`: the relay is not entitled to know why a command failed, and the
        /// phone needs to tell the user what to do about it.
        public var outcome: String?
        public var body: Data?

        public init(status: Int, outcome: String? = nil, body: Data? = nil) {
            self.status = status
            self.outcome = outcome
            self.body = body
        }
    }

    // MARK: - Canonical bytes

    /// The signing input, which is not JSON.
    ///
    /// Fixed field order, each field prefixed with its length as uint32be. The two integers carry
    /// an explicit `00000008` length too — Go's `writeUint64` delegates to `writeBytes`, so a
    /// version that emitted them bare would produce a byte string that looks identical in a hex
    /// dump's payload region and still fails to verify.
    ///
    /// `ek` and `p` contribute their raw decoded bytes, never their base64 text, so a relay cannot
    /// re-encode them differently.
    public static func canonical(_ envelope: Envelope, response: Bool) throws -> Data {
        guard envelope.v == version else { throw Failure.unsupportedVersion(envelope.v) }
        guard envelope.alg == suite else { throw Failure.unsupportedSuite(envelope.alg) }
        guard !envelope.node.trimmed.isEmpty, !envelope.req.trimmed.isEmpty else {
            throw Failure.missingIdentityField
        }
        guard envelope.ek.count == keyBytes else {
            throw Failure.badKeyLength(field: "ek", got: envelope.ek.count)
        }
        if !response {
            guard let app = envelope.app, !app.trimmed.isEmpty else { throw Failure.missingIdentityField }
        }

        var out = Data()
        out.append(contentsOf: Array((response ? domainResponse : domainRequest).utf8))
        out.append(0)
        appendLengthPrefixed(&out, uint64be(envelope.v))
        appendLengthPrefixed(&out, Data(envelope.node.utf8))
        appendLengthPrefixed(&out, Data(envelope.req.utf8))
        if !response, let app = envelope.app {
            appendLengthPrefixed(&out, Data(app.utf8))
        }
        appendLengthPrefixed(&out, uint64be(envelope.ts))
        appendLengthPrefixed(&out, envelope.ek)
        appendLengthPrefixed(&out, envelope.p)
        return out
    }

    private static func appendLengthPrefixed(_ out: inout Data, _ value: Data) {
        var length = UInt32(value.count).bigEndian
        withUnsafeBytes(of: &length) { out.append(contentsOf: $0) }
        out.append(value)
    }

    private static func uint64be(_ value: UInt64) -> Data {
        var big = value.bigEndian
        return withUnsafeBytes(of: &big) { Data($0) }
    }

    // MARK: - Signing

    public static func signRequest(_ envelope: inout Envelope, privateKey: Curve25519.Signing.PrivateKey) throws {
        envelope.s = try privateKey.signature(for: canonical(envelope, response: false))
    }

    public static func verifyRequest(_ envelope: Envelope, publicKey: Curve25519.Signing.PublicKey) throws {
        try verify(envelope, publicKey: publicKey, response: false)
    }

    public static func verifyResponse(_ envelope: Envelope, publicKey: Curve25519.Signing.PublicKey) throws {
        try verify(envelope, publicKey: publicKey, response: true)
    }

    private static func verify(
        _ envelope: Envelope,
        publicKey: Curve25519.Signing.PublicKey,
        response: Bool
    ) throws {
        let canonicalBytes = try canonical(envelope, response: response)
        guard let signature = envelope.s, signature.count == 64 else {
            throw Failure.signatureDoesNotVerify
        }
        guard publicKey.isValidSignature(signature, for: canonicalBytes) else {
            throw Failure.signatureDoesNotVerify
        }
    }

    // MARK: - Key agreement and AEAD

    /// HKDF-SHA256, extract then expand.
    public static func hkdf(secret: Data, salt: Data, info: Data, length: Int) throws -> Data {
        guard length > 0, length <= 255 * SHA256.byteCount else {
            throw Failure.badKeyLength(field: "hkdf", got: length)
        }
        let key = HKDF<SHA256>.deriveKey(
            inputKeyMaterial: SymmetricKey(data: secret),
            salt: salt,
            info: info,
            outputByteCount: length
        )
        return key.withUnsafeBytes { Data($0) }
    }

    /// X25519, rejecting the all-zero output that a peer key chosen to force a known shared secret
    /// would produce.
    public static func sharedSecret(
        privateKey: Curve25519.KeyAgreement.PrivateKey,
        peerPublic: Data
    ) throws -> Data {
        guard peerPublic.count == keyBytes else {
            throw Failure.badKeyLength(field: "peerPublic", got: peerPublic.count)
        }
        guard let publicKey = try? Curve25519.KeyAgreement.PublicKey(rawRepresentation: peerPublic) else {
            throw Failure.peerKeyUnusable
        }
        do {
            let secret = try privateKey.sharedSecretFromKeyAgreement(with: publicKey)
            return secret.withUnsafeBytes { Data($0) }
        } catch {
            throw Failure.peerKeyUnusable
        }
    }

    /// The AEAD key and IV for one direction of one request.
    ///
    /// The salt binds the node id, so a shared secret that somehow repeats across nodes still
    /// yields different keys. The IV is derived from the request id and the direction rather than
    /// being a fixed zero block: a fixed IV is only safe while every message uses a fresh key, and
    /// that invariant is exactly what a later "reuse the session key for a batch" optimisation
    /// would quietly break.
    public static func deriveKey(
        shared: Data,
        nodeID: String,
        requestID: String,
        response: Bool
    ) throws -> (key: SymmetricKey, iv: Data) {
        guard shared.count == keyBytes else {
            throw Failure.badKeyLength(field: "shared", got: shared.count)
        }
        let salt = Data("\(label):\(nodeID):\(requestID)".utf8)
        let keyBytes = try hkdf(
            secret: shared,
            salt: salt,
            info: Data((response ? "\(label)-resp" : "\(label)-req").utf8),
            length: self.keyBytes
        )
        let iv = try hkdf(
            secret: shared,
            salt: Data(requestID.utf8),
            info: Data((response ? "iv-resp" : "iv-req").utf8),
            length: ivBytes
        )
        return (SymmetricKey(data: keyBytes), iv)
    }

    /// Encrypts with AES-256-GCM. The tag is appended to the ciphertext, matching Node's
    /// `cipher.getAuthTag()` concatenation and Go's `gcm.Seal`.
    public static func seal(_ plaintext: Data, key: SymmetricKey, iv: Data) throws -> Data {
        guard iv.count == ivBytes else { throw Failure.badKeyLength(field: "iv", got: iv.count) }
        let nonce = try AES.GCM.Nonce(data: iv)
        let box = try AES.GCM.seal(plaintext, using: key, nonce: nonce)
        return box.ciphertext + box.tag
    }

    /// Decrypts and authenticates. A failure is terminal for the request, not a retry: a wrong
    /// key, a torn payload and a forged payload are indistinguishable here, and all three mean do
    /// not act.
    public static func open(_ ciphertext: Data, key: SymmetricKey, iv: Data) throws -> Data {
        guard iv.count == ivBytes else { throw Failure.badKeyLength(field: "iv", got: iv.count) }
        guard ciphertext.count >= tagBytes else { throw Failure.ciphertextFailedToAuthenticate }
        let nonce = try AES.GCM.Nonce(data: iv)
        let body = ciphertext.prefix(ciphertext.count - tagBytes)
        let tag = ciphertext.suffix(tagBytes)
        let box = try AES.GCM.SealedBox(nonce: nonce, ciphertext: body, tag: tag)
        do {
            return try AES.GCM.open(box, using: key)
        } catch {
            throw Failure.ciphertextFailedToAuthenticate
        }
    }

    // MARK: - Confirmation intent

    /// sha256(method | path | canonical body): the value a phone sends to prove a human confirmed
    /// this exact write.
    ///
    /// The agent recomputes it from the request it is about to forward, so a relay cannot edit the
    /// recipient of an SMS in flight and keep the confirmation valid. The body is canonicalised
    /// the same way the web console does it, which is why `CanonicalJSON` exists as a shared,
    /// test-vectored implementation rather than a call to `JSONSerialization`.
    public static func intentDigest(method: String, path: String, body: Data?) throws -> Data {
        let canonicalBody = try CanonicalJSON.canonicalize(body ?? Data())
        var input = Data()
        input.append(contentsOf: Array(method.trimmingCharacters(in: .whitespaces).uppercased().utf8))
        input.append(UInt8(ascii: "|"))
        input.append(contentsOf: Array(path.utf8))
        input.append(UInt8(ascii: "|"))
        input.append(canonicalBody)
        return Data(SHA256.hash(data: input))
    }
}

private extension String {
    var trimmed: String { trimmingCharacters(in: .whitespacesAndNewlines) }
}