import XCTest
import CryptoKit
@testable import Closer

final class InvitePayloadTests: XCTestCase {
    /// Tests the real `FirestoreService` shape end-to-end using a mock server.
    /// The inviter's key is stored locally under the invite code; the acceptor
    /// decrypts the phrase and unwraps the key; both keys must match.
    func testCreateAcceptRoundTripRecoversSameKey() async throws {
        let mockInvites = MockFirestoreInvites()
        let inviterKeyStore = InMemoryCoupleKeyStore()
        let acceptorKeyStore = InMemoryCoupleKeyStore()

        let inviterVM = PairingViewModel(invites: mockInvites, keyStore: inviterKeyStore)
        let acceptorVM = PairingViewModel(invites: mockInvites, keyStore: acceptorKeyStore)

        // Inviter creates an invite.
        let invite = try await inviterVM.startCreateInvite(uid: "inviter-uid")
        XCTAssertEqual(invite.code.count, 6)
        XCTAssertEqual(invite.recoveryPhrase.split(separator: " ").count, 10)

        guard let inviterKey = try inviterKeyStore.loadCoupleKey(for: invite.code) else {
            XCTFail("Inviter key missing")
            return
        }

        // Encrypt the inviter's key and phrase exactly as the real server would store them.
        let wrapped = try CoupleEncryptionManager.wrap(inviterKey, with: invite.recoveryPhrase)
        let encryptedPhrase = try CoupleEncryptionManager.encryptRecoveryPhrase(
            invite.recoveryPhrase,
            with: invite.code
        )
        let realPayload = InvitePayload(
            code: invite.code,
            wrappedCoupleKey: wrapped.ciphertext.base64EncodedString(),
            kdfSalt: wrapped.kdfSalt.base64EncodedString(),
            kdfParams: wrapped.kdfParams,
            encryptedRecoveryPhrase: encryptedPhrase
        )
        // Seed the mock with a placeholder, then replace with the real encrypted payload.
        _ = try await mockInvites.createInvite(uid: "inviter-uid", code: invite.code, recoveryPhrase: invite.recoveryPhrase)
        mockInvites.replaceStoredPayload(code: invite.code, payload: realPayload)

        // Acceptor accepts using the same code + phrase.
        try await acceptorVM.acceptInvite(code: invite.code, phrase: invite.recoveryPhrase)

        guard let acceptorKey = try acceptorKeyStore.loadCoupleKey(for: mockInvites.nextCoupleId) else {
            XCTFail("Acceptor key missing")
            return
        }
        XCTAssertEqual(acceptorKey.rawKey.bytes, inviterKey.rawKey.bytes)
    }

    /// A mismatched recovery phrase should fail during acceptance.
    func testWrongRecoveryPhraseRejects() async throws {
        let mockInvites = MockFirestoreInvites()
        let keyStore = InMemoryCoupleKeyStore()
        let vm = PairingViewModel(invites: mockInvites, keyStore: keyStore)

        let invite = try await vm.startCreateInvite(uid: "uid")

        let wrongPhrase = invite.recoveryPhrase + " wrong"
        do {
            try await vm.acceptInvite(code: invite.code, phrase: wrongPhrase)
            XCTFail("Expected phrase mismatch error")
        } catch {
            // Expected.
        }
    }
}

extension MockFirestoreInvites {
    fileprivate func replaceStoredPayload(code: String, payload: InvitePayload) {
        lock.lock()
        defer { lock.unlock() }
        invites[code] = payload
    }
}