import CryptoKit
import XCTest

@testable import MultiSimAdmin

/// Recomputes `contracts/test-vectors.json` in Swift.
///
/// The Go agent, the Node reference implementation and this app share no code, so agreement here is
/// evidence rather than a tautology: three encoders producing the same bytes for the same inputs is
/// what makes "the phone signs, the agent verifies" true across a language boundary. A failure in
/// this file means the encoding drifted, not that this side is broken.
final class RelayWireTests: XCTestCase {
    /// The vector file lives in the relay repository. Both paths are tried so the test works from
    /// the app package and from a checkout where the two repositories sit side by side; when
    /// neither exists the test is skipped rather than failed, because a missing sibling repository
    /// is not an encoding bug.
    private static let candidatePaths = [
        ProcessInfo.processInfo.environment["MSA_RELAY_VECTORS"],
        "/Users/chick/.local/opt/simadmin-relay/contracts/test-vectors.json",
        FileManager.default.currentDirectoryPath + "/../../simadmin-relay/contracts/test-vectors.json",
        FileManager.default.currentDirectoryPath + "/../../../simadmin-relay/contracts/test-vectors.json",
    ].compactMap { $0 }

    private func loadVectors() throws -> [String: Any] {
        for path in Self.candidatePaths where FileManager.default.fileExists(atPath: path) {
            let data = try Data(contentsOf: URL(fileURLWithPath: path))
            guard let object = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
                throw XCTSkip("vectors at \(path) are not a JSON object")
            }
            return object
        }
        throw XCTSkip("simadmin-relay/contracts/test-vectors.json not found; set MSA_RELAY_VECTORS")
    }

    private func hex(_ data: Data) -> String {
        data.map { String(format: "%02x", $0) }.joined()
    }

    private func unhex(_ text: String) -> Data {
        var out = Data()
        var index = text.startIndex
        while index < text.endIndex {
            guard let next = text.index(index, offsetBy: 2, limitedBy: text.endIndex),
                  let byte = UInt8(text[index..<next], radix: 16)
            else {
                XCTFail("not hex: \(text)")
                return Data()
            }
            out.append(byte)
            index = next
        }
        return out
    }

    /// The envelope in the vector file is JSON, so its byte fields are base64 there while the
    /// key-material sections are hex. Both decoders are needed and mixing them up produces an empty
    /// Data that fails as a length error three assertions later.
    private func unbase64(_ text: String) -> Data {
        guard let data = Data(base64Encoded: text) else {
            XCTFail("not base64: \(text)")
            return Data()
        }
        return data
    }

    /// RFC 8032 seed halves, which are the 32-byte raw representations CryptoKit wants. The vector
    /// file stores the full 64-byte expanded form that Go's ed25519 package uses.
    private func seed(_ expanded: String) -> Data {
        unhex(String(expanded.prefix(64)))
    }

    // MARK: - Envelope encoding

    func testCanonicalBytesAndSignaturesMatchTheGoAgent() throws {
        let vectors = try loadVectors()
        guard let canonical = vectors["canonical"] as? [String: Any],
              let input = canonical["input"] as? [String: Any],
              let bytes = canonical["bytes"] as? [String: String],
              let signatures = canonical["signature"] as? [String: String],
              let publicKeyHex = canonical["publicKey"] as? String,
              let responsePublicKeyHex = canonical["responsePublicKey"] as? String
        else {
            return XCTFail("the canonical section is missing or malformed")
        }

        let envelope = RelayWire.Envelope(
            v: input["v"] as? UInt64 ?? 0,
            node: input["node"] as? String ?? "",
            req: input["req"] as? String ?? "",
            app: input["app"] as? String,
            ts: input["ts"] as? UInt64 ?? 0,
            ek: unbase64(input["ek"] as? String ?? ""),
            alg: input["alg"] as? String ?? "",
            p: unbase64(input["p"] as? String ?? "")
        )

        XCTAssertEqual(hex(try RelayWire.canonical(envelope, response: false)), bytes["request"])
        XCTAssertEqual(hex(try RelayWire.canonical(envelope, response: true)), bytes["response"])

        let appSigningKey = try Curve25519.Signing.PrivateKey(rawRepresentation: seed(privateKeyHex))
        let agentSigningKey = try Curve25519.Signing.PrivateKey(rawRepresentation: seed(agentSignKeyHex))

        // The public halves published in the vector file must be the ones these private keys derive,
        // or the vector proves nothing about key ownership.
        XCTAssertEqual(hex(appSigningKey.publicKey.rawRepresentation), publicKeyHex)
        XCTAssertEqual(hex(agentSigningKey.publicKey.rawRepresentation), responsePublicKeyHex)

        // CryptoKit's Ed25519 signing is randomised, unlike Go's and Node's deterministic RFC 8032
        // implementation: signing the same bytes twice yields two different, equally valid
        // signatures. So the cross-language check here cannot be "the bytes are equal" — it is the
        // stronger property that actually matters, which is that each side verifies what the other
        // produced. Both directions are asserted, because a phone whose signatures the agent
        // rejected would be a total failure that byte-equality would never have caught either.
        try RelayWire.verifyRequest(
            envelopeWithSignature(envelope, unhex(signatures["request"] ?? ""), app: true),
            publicKey: appSigningKey.publicKey
        )
        try RelayWire.verifyResponse(
            envelopeWithSignature(envelope, unhex(signatures["response"] ?? ""), app: false),
            publicKey: agentSigningKey.publicKey
        )

        var signedRequest = envelope
        try RelayWire.signRequest(&signedRequest, privateKey: appSigningKey)
        try RelayWire.verifyRequest(signedRequest, publicKey: appSigningKey.publicKey)
        // And the Go side must accept what this side signed: the vectors cannot prove that from
        // here, so it is asserted by the agent's own test suite against a signature captured from
        // a device. What can be checked locally is that the signing input is the vector's bytes.
        XCTAssertEqual(
            hex(try RelayWire.canonical(signedRequest, response: false)),
            bytes["request"],
            "the phone must sign exactly the bytes the agent verifies"
        )

        // Domain separation: a request signature must not verify as a response signature.
        XCTAssertThrowsError(try RelayWire.verifyResponse(signedRequest, publicKey: appSigningKey.publicKey))
    }

    /// Rebuilds an envelope carrying a signature produced elsewhere, since the response form has no
    /// `app` field and the request form does.
    private func envelopeWithSignature(
        _ envelope: RelayWire.Envelope,
        _ signature: Data,
        app: Bool
    ) -> RelayWire.Envelope {
        var copy = envelope
        copy.app = app ? envelope.app : nil
        copy.s = signature
        return copy
    }

    /// RFC 8032 test vector private keys, repeated here so the test does not depend on the vector
    /// file to know whose key is whose.
    private let privateKeyHex =
        "9d61b19deffd5a60ba844af492ec2cc44449c5697b326919703bac031cae7f60"
        + "d75a980182b10ab7d54bfed3c964073a0ee172f3daa62325af021a68f707511a"
    private let agentSignKeyHex =
        "4ccd089b28ff96da9db6c346ec114e0f5b8a319f35aba624da8cf6ed4fb8a6fb"
        + "3d4017c3e843895a92b70aa74d1b7ebc9c982ccf2ec4968cc0cd55f12af4660c"

    func testAOneByteChangeBreaksTheSignature() throws {
        let envelope = RelayWire.Envelope(
            node: "nd_7f3a9c2e",
            req: "3f2504e0-4f89-41d3-9a0c-0305e82c3301",
            app: "ap_5b1d",
            ts: 1_727_289_600,
            ek: Data(repeating: 7, count: 32),
            p: Data([0, 1, 2, 0xFE, 0xFF])
        )
        let key = try Curve25519.Signing.PrivateKey(rawRepresentation: seed(privateKeyHex))
        var signed = envelope
        try RelayWire.signRequest(&signed, privateKey: key)
        try RelayWire.verifyRequest(signed, publicKey: key.publicKey)

        // Retargeting the node is what a compromised relay would try first.
        var retargeted = signed
        retargeted.node = "nd_other001"
        XCTAssertThrowsError(try RelayWire.verifyRequest(retargeted, publicKey: key.publicKey)) { error in
            XCTAssertEqual(error as? RelayWire.Failure, .signatureDoesNotVerify)
        }
        // Swapping the ciphertext is the second thing it would try.
        var edited = signed
        edited.p = Data([9, 9, 9])
        XCTAssertThrowsError(try RelayWire.verifyRequest(edited, publicKey: key.publicKey))
        // The ephemeral public key is covered, so a substituted one cannot be hidden.
        var swapped = signed
        swapped.ek = Data(repeating: 3, count: 32)
        XCTAssertThrowsError(try RelayWire.verifyRequest(swapped, publicKey: key.publicKey))
    }

    // MARK: - Key derivation and AEAD

    func testKeyDerivationAndAEADMatchTheGoAgent() throws {
        let vectors = try loadVectors()
        guard let x25519 = vectors["x25519"] as? [String: String],
              let aead = vectors["aead"] as? [String: String]
        else {
            return XCTFail("the x25519 or aead section is missing")
        }

        let privateKey = try Curve25519.KeyAgreement.PrivateKey(
            rawRepresentation: unhex(x25519["appPrivate"] ?? "")
        )
        let shared = try RelayWire.sharedSecret(
            privateKey: privateKey,
            peerPublic: unhex(x25519["agentPublic"] ?? "")
        )
        XCTAssertEqual(hex(shared), x25519["sharedSecret"])

        let request = try RelayWire.deriveKey(
            shared: shared,
            nodeID: "nd_7f3a9c2e",
            requestID: "3f2504e0-4f89-41d3-9a0c-0305e82c3301",
            response: false
        )
        XCTAssertEqual(hex(request.key.withUnsafeBytes { Data($0) }), x25519["requestKey"])
        XCTAssertEqual(hex(request.iv), x25519["requestIV"])

        let response = try RelayWire.deriveKey(
            shared: shared,
            nodeID: "nd_7f3a9c2e",
            requestID: "3f2504e0-4f89-41d3-9a0c-0305e82c3301",
            response: true
        )
        XCTAssertEqual(hex(response.key.withUnsafeBytes { Data($0) }), x25519["responseKey"])
        XCTAssertEqual(hex(response.iv), x25519["responseIV"])
        // Direction separation: the same shared secret must not produce the same key or nonce.
        XCTAssertNotEqual(x25519["requestKey"], x25519["responseKey"])
        XCTAssertNotEqual(x25519["requestIV"], x25519["responseIV"])

        let sealed = try RelayWire.seal(
            unhex(aead["plaintext"] ?? ""),
            key: request.key,
            iv: request.iv
        )
        XCTAssertEqual(hex(sealed), aead["ciphertext"])
        let opened = try RelayWire.open(sealed, key: request.key, iv: request.iv)
        XCTAssertEqual(hex(opened), aead["plaintext"])
    }

    func testHKDFMatchesRFC5869() throws {
        let vectors = try loadVectors()
        guard let hkdf = vectors["hkdf"] as? [String: Any] else {
            return XCTFail("the hkdf section is missing")
        }
        let output = try RelayWire.hkdf(
            secret: unhex(hkdf["secret"] as? String ?? ""),
            salt: unhex(hkdf["salt"] as? String ?? ""),
            info: unhex(hkdf["info"] as? String ?? ""),
            length: hkdf["length"] as? Int ?? 0
        )
        XCTAssertEqual(hex(output), hkdf["output"] as? String)
    }

    /// A tampered ciphertext must fail to authenticate rather than decrypt to something.
    func testATamperedCiphertextDoesNotOpen() throws {
        let key = SymmetricKey(data: Data(repeating: 1, count: 32))
        let iv = Data(repeating: 2, count: 12)
        var sealed = try RelayWire.seal(Data("hello".utf8), key: key, iv: iv)
        sealed[sealed.count - 1] ^= 0x01
        XCTAssertThrowsError(try RelayWire.open(sealed, key: key, iv: iv)) { error in
            XCTAssertEqual(error as? RelayWire.Failure, .ciphertextFailedToAuthenticate)
        }
    }

    // MARK: - Intent digest

    func testIntentDigestMatchesTheGoAgent() throws {
        let vectors = try loadVectors()
        guard let intent = vectors["intent"] as? [String: String] else {
            return XCTFail("the intent section is missing")
        }
        let digest = try RelayWire.intentDigest(
            method: intent["method"] ?? "",
            path: intent["path"] ?? "",
            body: Data((intent["body"] ?? "").utf8)
        )
        XCTAssertEqual(hex(digest), intent["digest"])

        // Formatting must not matter; content must.
        let reordered = try RelayWire.intentDigest(
            method: "post",
            path: "/api/sms/send",
            body: Data(#"{"content":"hello","phone_number":"+8613800000000"}"#.utf8)
        )
        XCTAssertEqual(hex(reordered), intent["digest"])
        let edited = try RelayWire.intentDigest(
            method: "POST",
            path: "/api/sms/send",
            body: Data(#"{"phone_number":"+8613800000001","content":"hello"}"#.utf8)
        )
        XCTAssertNotEqual(hex(edited), intent["digest"])
    }

    // MARK: - Canonical JSON

    func testCanonicalJSONMatchesTheGoAgent() throws {
        let vectors = try loadVectors()
        guard let pairs = vectors["canonicalJson"] as? [[String: String]] else {
            return XCTFail("the canonicalJson section is missing")
        }
        for pair in pairs {
            let input = pair["input"] ?? ""
            let expected = pair["output"] ?? ""
            let got = try CanonicalJSON.canonicalize(input)
            XCTAssertEqual(String(decoding: got, as: UTF8.self), expected, "canonical JSON for \(input)")
        }
    }

    func testCanonicalJSONRejectsEveryInputTheVectorsMarkInvalid() throws {
        let vectors = try loadVectors()
        guard let rejected = vectors["canonicalJsonRejected"] as? [String] else {
            return XCTFail("the canonicalJsonRejected section is missing")
        }
        for input in rejected {
            XCTAssertThrowsError(try CanonicalJSON.canonicalize(input), "accepted \(input)")
        }
    }

    func testCanonicalJSONRejectsWhatTheProtocolWillNotHash() {
        for input in ["{", #"{"a":1} trailing"#, #"{"a":+1}"#, #"{"a":.5}"#, #"{"a":NaN}"#,
                      #"{"a":1.0}"#, #"{"a":1e3}"#, #"{"a":01}"#, #"{"a":-0}"#] {
            XCTAssertThrowsError(try CanonicalJSON.canonicalize(input), "accepted \(input)")
        }
        XCTAssertEqual(try CanonicalJSON.canonicalize(""), Data())
        XCTAssertEqual(try CanonicalJSON.canonicalize("   "), Data())
    }

    /// The three key-sort shortcuts disagree, so this pins the one the protocol uses: UTF-8 byte
    /// order. Swift's default `String` comparison uses Unicode collation, which would put "é"
    /// before "z"; Foundation's `.sortedKeys` does the same.
    func testCanonicalJSONSortsKeysByUTF8Bytes() throws {
        XCTAssertEqual(
            String(decoding: try CanonicalJSON.canonicalize(#"{"é":1,"z":2}"#), as: UTF8.self),
            #"{"z":2,"é":1}"#
        )
        XCTAssertEqual(
            String(decoding: try CanonicalJSON.canonicalize("{\"\u{10000}\":1,\"\u{FFFD}\":2}"), as: UTF8.self),
            "{\"\u{FFFD}\":2,\"\u{10000}\":1}"
        )
    }
}