import XCTest
import CryptoKit
@testable import Closer

final class KeyboxCryptoTests: XCTestCase {
    private let info = Data("couple-1|q-1|sender-1|recipient-1".utf8)

    func testWrapUnwrapRoundTrip() throws {
        let recipient = P256.KeyAgreement.PrivateKey()
        let plaintext = Data("secret one-time answer key".utf8)

        let keybox = try KeyboxCrypto.wrap(
            plaintext: plaintext,
            recipientPublicKey: recipient.publicKey,
            info: info
        )
        XCTAssertEqual(keybox.ephemeralPublicKey.count, 65)
        XCTAssertTrue(keybox.ephemeralPublicKey.first == 0x04)
        XCTAssertEqual(keybox.mac.count, 32)

        let decoded = try KeyboxCrypto.encode(keybox)
        XCTAssertTrue(decoded.hasPrefix(KeyboxCrypto.keyboxPrefix))

        let reloaded = try KeyboxCrypto.decode(decoded)
        let recovered = try KeyboxCrypto.unwrap(reloaded, recipientPrivateKey: recipient, info: info)
        XCTAssertEqual(recovered, plaintext)
    }

    func testAADMismatchRejects() throws {
        let recipient = P256.KeyAgreement.PrivateKey()
        let plaintext = Data("aad-bound secret".utf8)
        let keybox = try KeyboxCrypto.wrap(
            plaintext: plaintext,
            recipientPublicKey: recipient.publicKey,
            info: info
        )
        let wrongInfo = Data("different info".utf8)
        XCTAssertThrowsError(try KeyboxCrypto.unwrap(keybox, recipientPrivateKey: recipient, info: wrongInfo))
    }

    func testTamperedMACRejects() throws {
        let recipient = P256.KeyAgreement.PrivateKey()
        let plaintext = Data("mac-bound secret".utf8)
        var keybox = try KeyboxCrypto.wrap(
            plaintext: plaintext,
            recipientPublicKey: recipient.publicKey,
            info: info
        )
        var macBytes = Array(keybox.mac)
        macBytes[0] ^= 0x01
        keybox = Keybox(
            ephemeralPublicKey: keybox.ephemeralPublicKey,
            ciphertext: keybox.ciphertext,
            mac: Data(macBytes)
        )
        XCTAssertThrowsError(try KeyboxCrypto.unwrap(keybox, recipientPrivateKey: recipient, info: info))
    }

    func testTamperedCiphertextRejects() throws {
        let recipient = P256.KeyAgreement.PrivateKey()
        let plaintext = Data("ct-bound secret".utf8)
        var keybox = try KeyboxCrypto.wrap(
            plaintext: plaintext,
            recipientPublicKey: recipient.publicKey,
            info: info
        )
        var ctBytes = Array(keybox.ciphertext)
        ctBytes[12 + 3] ^= 0x01
        keybox = Keybox(
            ephemeralPublicKey: keybox.ephemeralPublicKey,
            ciphertext: Data(ctBytes),
            mac: keybox.mac
        )
        XCTAssertThrowsError(try KeyboxCrypto.unwrap(keybox, recipientPrivateKey: recipient, info: info))
    }

    func testInfoStringMismatchRejects() throws {
        let recipient = P256.KeyAgreement.PrivateKey()
        let plaintext = Data("info-bound secret".utf8)
        let keybox = try KeyboxCrypto.wrap(
            plaintext: plaintext,
            recipientPublicKey: recipient.publicKey,
            info: Data("a|b|c|d".utf8)
        )
        XCTAssertThrowsError(
            try KeyboxCrypto.unwrap(
                keybox,
                recipientPrivateKey: recipient,
                info: Data("a|b|c|e".utf8)
            )
        )
    }

    func testEncodeDecodePreservesFields() throws {
        let recipient = P256.KeyAgreement.PrivateKey()
        let keybox = try KeyboxCrypto.wrap(
            plaintext: Data("encode decode".utf8),
            recipientPublicKey: recipient.publicKey,
            info: info
        )
        let encoded = try KeyboxCrypto.encode(keybox)
        let decoded = try KeyboxCrypto.decode(encoded)
        XCTAssertEqual(decoded.ephemeralPublicKey, keybox.ephemeralPublicKey)
        XCTAssertEqual(decoded.ciphertext, keybox.ciphertext)
        XCTAssertEqual(decoded.mac, keybox.mac)
    }
}