Closer/iphone/Closer/Services/FirestoreService.swift

import Foundation
import FirebaseFirestore
import FirebaseAuth
import FirebaseFunctions

// MARK: - E2EE callable contract (Batch 2)
//
// `createInviteCallable` (functions/src/couples/createInviteCallable.ts) requires
// a strict-E2EE payload in this exact shape:
//
//   {
//     code: String,              // 6-char Crockford alphabet [A-HJ-NP-Z2-9]
//     wrappedCoupleKey: String,  // base64( AES-256-GCM( keyMaterial, KEK=Argon2id(phrase, salt), aad="closer_couple_key" ) )
//     kdfSalt: String,           // base64( 16 random bytes )
//     kdfParams: String,         // "argon2id;v=19;m=47104;t=3;p=1"
//     encryptedRecoveryPhrase: String  // base64( salt[16] || AES-256-GCM(phrase, KEK=Argon2id(code, salt), aad="closer_invite_phrase") )
//   }
//
// `acceptInviteCallable` returns the same four E2EE fields (plus coupleId and
// inviterUserId). The acceptor must:
//   1. Derive the phrase-decryption KEK from the invite code and the embedded salt.
//   2. Decrypt `encryptedRecoveryPhrase` with AAD "closer_invite_phrase".
//   3. Derive the couple-key KEK from the recovery phrase and `kdfSalt`.
//   4. Unwrap `wrappedCoupleKey` with AAD "closer_couple_key".
//   5. Store the unwrapped key in the Keychain via CoupleKeyStore.
//
// Field order in the callable dictionary is not semantically meaningful for JSON,
// but the values above must all be present and non-nil.

// MARK: - FirestoreInvitesProtocol

/// Protocol covering the invite create/accept callables. Allows the production
/// `FirestoreService` and a deterministic test fake to share the same contract.
public protocol FirestoreInvitesProtocol: Sendable {
    func createInvite(uid: String, code: String, recoveryPhrase: String) async throws -> InvitePayload
    func acceptInvite(code: String, inviterUserId: String, recoveryPhrase: String) async throws -> AcceptResult
}

// MARK: - Invite Payload Models

public struct InvitePayload: Sendable {
    public let code: String
    public let wrappedCoupleKey: String
    public let kdfSalt: String
    public let kdfParams: String
    public let encryptedRecoveryPhrase: String

    public init(
        code: String,
        wrappedCoupleKey: String,
        kdfSalt: String,
        kdfParams: String,
        encryptedRecoveryPhrase: String
    ) {
        self.code = code
        self.wrappedCoupleKey = wrappedCoupleKey
        self.kdfSalt = kdfSalt
        self.kdfParams = kdfParams
        self.encryptedRecoveryPhrase = encryptedRecoveryPhrase
    }

    var callableDictionary: [String: Any] {
        [
            "code": code,
            "wrappedCoupleKey": wrappedCoupleKey,
            "kdfSalt": kdfSalt,
            "kdfParams": kdfParams,
            "encryptedRecoveryPhrase": encryptedRecoveryPhrase
        ]
    }
}

public struct AcceptResult: Sendable {
    public let coupleId: String
    public let inviterUserId: String
    public let wrappedCoupleKey: String
    public let kdfSalt: String
    public let kdfParams: String
    public let encryptedRecoveryPhrase: String

    public init(
        coupleId: String,
        inviterUserId: String,
        wrappedCoupleKey: String,
        kdfSalt: String,
        kdfParams: String,
        encryptedRecoveryPhrase: String
    ) {
        self.coupleId = coupleId
        self.inviterUserId = inviterUserId
        self.wrappedCoupleKey = wrappedCoupleKey
        self.kdfSalt = kdfSalt
        self.kdfParams = kdfParams
        self.encryptedRecoveryPhrase = encryptedRecoveryPhrase
    }
}

// MARK: - Firestore Service

final class FirestoreService: @unchecked Sendable {
    static let shared = FirestoreService()

    let db = Firestore.firestore()
    let functions = Functions.functions()

    private init() {}

    // MARK: - Collection References

    func usersCollection() -> CollectionReference {
        db.collection("users")
    }

    func userDocument(_ userId: String) -> DocumentReference {
        usersCollection().document(userId)
    }

    func couplesCollection() -> CollectionReference {
        db.collection("couples")
    }

    func coupleDocument(_ coupleId: String) -> DocumentReference {
        couplesCollection().document(coupleId)
    }

    func invitesCollection() -> CollectionReference {
        db.collection("invites")
    }

    func inviteDocument(_ code: String) -> DocumentReference {
        invitesCollection().document(code)
    }

    // MARK: - Subcollections

    func dailyQuestionRef(coupleId: String, date: String) -> DocumentReference {
        coupleDocument(coupleId)
            .collection("daily_question")
            .document(date)
    }

    func dailyAnswerRef(coupleId: String, date: String, userId: String) -> DocumentReference {
        dailyQuestionRef(coupleId: coupleId, date: date)
            .collection("answers")
            .document(userId)
    }

    func questionThreadsRef(coupleId: String) -> CollectionReference {
        coupleDocument(coupleId).collection("question_threads")
    }

    func dateSwipesRef(coupleId: String) -> CollectionReference {
        coupleDocument(coupleId).collection("date_swipes")
    }

    func dateMatchesRef(coupleId: String) -> CollectionReference {
        coupleDocument(coupleId).collection("date_matches")
    }

    func bucketListRef(coupleId: String) -> CollectionReference {
        coupleDocument(coupleId).collection("bucket_list")
    }

    func capsulesRef(coupleId: String) -> CollectionReference {
        coupleDocument(coupleId).collection("capsules")
    }

    func sessionsRef(coupleId: String) -> CollectionReference {
        coupleDocument(coupleId).collection("sessions")
    }

    func entitlementDocument(_ userId: String) -> DocumentReference {
        userDocument(userId)
            .collection("entitlements")
            .document("premium")
    }

    func fcmTokensRef(_ userId: String) -> CollectionReference {
        userDocument(userId).collection("fcmTokens")
    }

    // MARK: - Helpers

    func userId() throws -> String {
        guard let uid = Auth.auth().currentUser?.uid else {
            throw FirestoreError.notAuthenticated
        }
        return uid
    }

    func setDocument<T: Encodable>(_ value: T, at document: DocumentReference, merge: Bool = true) async throws {
        if merge {
            try await document.setData(value.asDictionary(), merge: true)
        } else {
            try await document.setData(value.asDictionary())
        }
    }

    func getDocument<T: Decodable>(at document: DocumentReference) async throws -> T? {
        let snapshot = try await document.getDocument()
        guard snapshot.exists else { return nil }
        return try snapshot.data(as: T.self)
    }

    func getDocuments<T: Decodable>(in collection: CollectionReference) async throws -> [T] {
        let snapshot = try await collection.getDocuments()
        return try snapshot.documents.compactMap { try $0.data(as: T.self) }
    }

    func queryDocuments<T: Decodable>(
        in collection: CollectionReference,
        where field: String,
        isEqualTo value: Any
    ) async throws -> [T] {
        let snapshot = try await collection.whereField(field, isEqualTo: value).getDocuments()
        return try snapshot.documents.compactMap { try $0.data(as: T.self) }
    }
}

// MARK: - Callable Functions

extension FirestoreService: FirestoreInvitesProtocol {
    /// Creates a strict-E2EE invite. The caller (PairingViewModel) supplies the
    /// 6-character Crockford code and recovery phrase; this method generates the
    /// couple key, wraps it with the phrase, encrypts the phrase with the code,
    /// and forwards the full payload to `createInviteCallable`.
    public func createInvite(uid: String, code: String, recoveryPhrase: String) async throws -> InvitePayload {
        let key = try CoupleEncryptionManager.generateCoupleKey()
        let wrapped = try CoupleEncryptionManager.wrap(key, with: recoveryPhrase)
        let encryptedPhrase = try CoupleEncryptionManager.encryptRecoveryPhrase(recoveryPhrase, with: code)

        let payload = InvitePayload(
            code: code,
            wrappedCoupleKey: wrapped.ciphertext.base64EncodedString(),
            kdfSalt: wrapped.kdfSalt.base64EncodedString(),
            kdfParams: wrapped.kdfParams,
            encryptedRecoveryPhrase: encryptedPhrase
        )

        let result = try await functions.httpsCallable("createInviteCallable").call(payload.callableDictionary)
        guard let data = result.data as? [String: Any],
              let returnedCode = data["code"] as? String else {
            throw FirestoreError.invalidResponse
        }
        // The server may return a different code in theory; the contract says it returns
        // the same code we sent. Include it in the returned payload for the UI.
        return InvitePayload(
            code: returnedCode,
            wrappedCoupleKey: payload.wrappedCoupleKey,
            kdfSalt: payload.kdfSalt,
            kdfParams: payload.kdfParams,
            encryptedRecoveryPhrase: payload.encryptedRecoveryPhrase
        )
    }

    /// Accepts an invite and recovers the couple key. The acceptor already knows
    /// the invite code (they typed it) but must receive the E2EE fields from the
    /// server to decrypt the recovery phrase and unwrap the couple key.
    public func acceptInvite(code: String, inviterUserId: String, recoveryPhrase: String) async throws -> AcceptResult {
        let result = try await functions.httpsCallable("acceptInviteCallable").call(["code": code])
        guard let data = result.data as? [String: Any],
              let coupleId = data["coupleId"] as? String,
              let inviter = data["inviterUserId"] as? String,
              let wrappedKey = data["wrappedCoupleKey"] as? String,
              let salt = data["kdfSalt"] as? String,
              let params = data["kdfParams"] as? String,
              let encryptedPhrase = data["encryptedRecoveryPhrase"] as? String else {
            throw FirestoreError.invalidResponse
        }
        return AcceptResult(
            coupleId: coupleId,
            inviterUserId: inviter,
            wrappedCoupleKey: wrappedKey,
            kdfSalt: salt,
            kdfParams: params,
            encryptedRecoveryPhrase: encryptedPhrase
        )
    }
}

// MARK: - Legacy callable wrappers (kept for existing call sites)

extension FirestoreService {
    func leaveCoupleCallable() async throws {
        let result = try await functions.httpsCallable("leaveCoupleCallable").call()
        guard let success = (result.data as? [String: Any])?["success"] as? Bool, success else {
            throw FirestoreError.invalidResponse
        }
    }

    func syncEntitlementCallable() async throws -> Entitlement {
        let result = try await functions.httpsCallable("syncEntitlement").call()
        guard let data = result.data as? [String: Any] else {
            throw FirestoreError.invalidResponse
        }
        return Entitlement(
            id: UUID().uuidString,
            userId: try userId(),
            source: "sync",
            productId: "closer_premium",
            isActive: data["premium"] as? Bool ?? false,
            expiresAt: (data["expiresAt"] as? Timestamp)?.dateValue(),
            updatedAt: Date()
        )
    }

    func sendGentleReminderCallable() async throws {
        try await functions.httpsCallable("sendGentleReminderCallable").call()
    }

    func deleteUserCallable() async throws {
        let result = try await functions.httpsCallable("deleteUserCallable").call()
        guard let success = (result.data as? [String: Any])?["success"] as? Bool, success else {
            throw FirestoreError.invalidResponse
        }
    }

    func updateUserCallable(displayName: String, bio: String?) async throws {
        var data: [String: Any] = ["displayName": displayName]
        if let bio = bio {
            data["bio"] = bio
        }
        let result = try await functions.httpsCallable("updateUserCallable").call(data)
        guard let success = (result.data as? [String: Any])?["success"] as? Bool, success else {
            throw FirestoreError.invalidResponse
        }
    }

    func exportDataCallable() async throws {
        let result = try await functions.httpsCallable("exportDataCallable").call()
        guard let success = (result.data as? [String: Any])?["success"] as? Bool, success else {
            throw FirestoreError.invalidResponse
        }
    }
}

// MARK: - Errors

enum FirestoreError: LocalizedError {
    case notAuthenticated
    case invalidResponse
    case documentNotFound
    case permissionDenied

    var errorDescription: String? {
        switch self {
        case .notAuthenticated: return "User is not signed in."
        case .invalidResponse: return "Invalid server response."
        case .documentNotFound: return "Document not found."
        case .permissionDenied: return "Permission denied."
        }
    }
}

// MARK: - Encodable Helpers

extension Encodable {
    func asDictionary() throws -> [String: Any] {
        let data = try JSONEncoder().encode(self)
        guard let dict = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            throw EncodingError.invalidValue(self, .init(codingPath: [], debugDescription: "Not a dictionary"))
        }
        return dict
    }
}