Closer/iphone/Closer/Pairing/PairingViewModel.swift

import Foundation
import SwiftUI

// MARK: - Pairing Errors

public enum PairingError: Error, LocalizedError {
    case notAuthenticated
    case missingCoupleKey
    case invalidInviteCode
    case recoveryPhraseMismatch
    case serverRejected(String)

    public var errorDescription: String? {
        switch self {
        case .notAuthenticated:
            return "You must be signed in to pair."
        case .missingCoupleKey:
            return "Couple encryption key is missing. Ask your partner to re-share the invite."
        case .invalidInviteCode:
            return "Invite code must be 6 characters."
        case .recoveryPhraseMismatch:
            return "The recovery phrase does not match this invite."
        case .serverRejected(let message):
            return message
        }
    }
}

// MARK: - Pairing State

public enum PairingState: Sendable {
    case idle
    case creatingInvite
    case acceptingInvite
    case paired(coupleId: String, partnerUserId: String)
    case failed(error: PairingError)
}

// MARK: - Invite Result

public struct CreatedInvite: Sendable {
    public let code: String
    public let recoveryPhrase: String
    public let expiresAt: Date?

    public init(code: String, recoveryPhrase: String, expiresAt: Date? = nil) {
        self.code = code
        self.recoveryPhrase = recoveryPhrase
        self.expiresAt = expiresAt
    }
}

// MARK: - Pairing View Model

/// Orchestrates E2EE invite creation and acceptance.
///
/// Responsibilities:
/// - Generate a 6-character Crockford invite code.
/// - Generate a fresh couple key + recovery phrase for the inviter.
/// - Store the inviter's couple key locally before the server call (so the
///   inviter can answer daily questions immediately after sharing the code).
/// - Forward the strict-E2EE payload to `FirestoreInvitesProtocol`.
/// - On acceptance, decrypt the recovery phrase with the invite code, unwrap the
///   couple key, and persist it via `CoupleKeyStoreProtocol`.
///
/// The ViewModel never logs or surfaces the raw key material; the recovery phrase
/// is returned to the UI only once (for display/copy), then it is the user's responsibility.
@MainActor
public final class PairingViewModel: ObservableObject {
    @Published public private(set) var state: PairingState = .idle

    private let invites: FirestoreInvitesProtocol
    private let keyStore: CoupleKeyStoreProtocol
    private let codeChars = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" // Crockford, no I/O/0/1

    public init(
        invites: FirestoreInvitesProtocol = FirestoreService.shared,
        keyStore: CoupleKeyStoreProtocol = CoupleKeyStore()
    ) {
        self.invites = invites
        self.keyStore = keyStore
    }

    /// Generates a new strict-E2EE invite.
    /// - Returns: the public invite code and the recovery phrase to display to the user.
    public func startCreateInvite(uid: String) async throws -> CreatedInvite {
        state = .creatingInvite
        defer {
            if case .creatingInvite = state {
                state = .idle
            }
        }

        let code = generateSixCharCode()
        let phrase = try RecoveryKeyManager.generatePhrase()
        let key = try CoupleEncryptionManager.generateCoupleKey()

        // Store the key locally keyed by the invite code. After the partner accepts,
        // `completePairing(coupleId:)` migrates it to the real coupleId.
        try keyStore.storeCoupleKey(key, for: code)

        let payload = try await invites.createInvite(
            uid: uid,
            code: code,
            recoveryPhrase: phrase
        )

        // If the server changed the code (not expected), re-key the stored key.
        if payload.code != code {
            try keyStore.deleteCoupleKey(for: code)
            try keyStore.storeCoupleKey(key, for: payload.code)
        }

        return CreatedInvite(
            code: payload.code,
            recoveryPhrase: phrase,
            expiresAt: nil
        )
    }

    /// Accepts an invite using the invite code and recovery phrase supplied by the partner.
    /// - Parameters:
    ///   - code: 6-character invite code.
    ///   - phrase: 10-word recovery phrase (normalized before use).
    public func acceptInvite(code: String, phrase: String) async throws {
        guard code.count == 6 else {
            throw PairingError.invalidInviteCode
        }
        let normalizedPhrase = RecoveryKeyManager.normalize(phrase)
        guard try RecoveryKeyManager.isWellFormed(normalizedPhrase) else {
            throw PairingError.recoveryPhraseMismatch
        }

        state = .acceptingInvite
        defer {
            if case .acceptingInvite = state {
                state = .idle
            }
        }

        let result = try await invites.acceptInvite(
            code: code,
            inviterUserId: "", // The server returns the real inviterUserId; this placeholder is unused.
            recoveryPhrase: normalizedPhrase
        )

        // The acceptor decrypts the recovery phrase using the invite code they typed.
        // This validates that the server returned the E2EE blob for the same code.
        let decryptedPhrase = try CoupleEncryptionManager.decryptRecoveryPhrase(
            result.encryptedRecoveryPhrase,
            with: code
        )
        guard RecoveryKeyManager.normalize(decryptedPhrase) == normalizedPhrase else {
            throw PairingError.recoveryPhraseMismatch
        }

        let wrapped = WrappedCoupleKey(
            ciphertext: Data(base64Encoded: result.wrappedCoupleKey) ?? Data(),
            kdfSalt: Data(base64Encoded: result.kdfSalt) ?? Data(),
            kdfParams: result.kdfParams
        )
        let key = try CoupleEncryptionManager.unwrap(wrapped, with: normalizedPhrase)
        try keyStore.storeCoupleKey(key, for: result.coupleId)

        state = .paired(coupleId: result.coupleId, partnerUserId: result.inviterUserId)
    }

    /// Migrates a locally-stored couple key from a temporary invite-code key to the
    /// permanent coupleId once the server confirms the couple was created.
    public func completePairing(inviteCode: String, coupleId: String) throws {
        guard let key = try keyStore.loadCoupleKey(for: inviteCode) else {
            throw PairingError.missingCoupleKey
        }
        try keyStore.storeCoupleKey(key, for: coupleId)
        try keyStore.deleteCoupleKey(for: inviteCode)
    }

    private func generateSixCharCode() -> String {
        var rng = SystemRandomNumberGenerator()
        return String((0..<6).map { _ in codeChars.randomElement(using: &rng)! })
    }
}