swift-api-client/Sources/MessagingAPIClient/Auth/KeychainAuthProvider.swift

import Foundation
import Security
import LilithLogging

// MARK: - Keychain Auth Provider

/// Keychain-backed authentication provider.
///
/// Stores access and refresh tokens securely in the iOS/macOS Keychain.
/// Handles SSO token exchange, JWT-based user extraction, and automatic
/// retry with exponential backoff on refresh failures.
public final class KeychainAuthProvider: AuthProvider, @unchecked Sendable {

    // MARK: - Configuration

    private let keychainService: String
    private let accessTokenKey = "access_token"
    private let refreshTokenKey = "refresh_token"
    private let baseURL: String

    // MARK: - State

    private var _isAuthenticated: Bool = false
    private var isRefreshing: Bool = false
    private let lock = NSLock()

    public var isAuthenticated: Bool {
        lock.lock()
        defer { lock.unlock() }
        return _isAuthenticated
    }

    public var accessToken: String? {
        readKeychain(key: accessTokenKey)
    }

    // MARK: - Initialization

    /// Create a Keychain-backed auth provider.
    /// - Parameters:
    ///   - keychainService: The Keychain service identifier (defaults to app bundle ID).
    ///   - baseURL: The base API URL for token exchange and refresh endpoints.
    public init(
        keychainService: String = Bundle.main.bundleIdentifier ?? "com.lilith.messenger",
        baseURL: String
    ) {
        self.keychainService = keychainService
        self.baseURL = baseURL

        // Restore authentication state from existing Keychain tokens
        if let token = readKeychain(key: accessTokenKey), isTokenValid(token) {
            lock.lock()
            _isAuthenticated = true
            lock.unlock()
        } else {
            clearCredentials()
        }
    }

    // MARK: - AuthProvider

    public func getAccessToken() async throws -> String {
        guard let token = accessToken else {
            throw AuthError.noRefreshToken
        }
        return token
    }

    public func refreshToken() async throws {
        lock.lock()
        guard !isRefreshing else {
            lock.unlock()
            return
        }
        isRefreshing = true
        lock.unlock()

        defer {
            lock.lock()
            isRefreshing = false
            lock.unlock()
        }

        guard let currentRefreshToken = readKeychain(key: refreshTokenKey) else {
            clearCredentials()
            throw AuthError.noRefreshToken
        }

        let url = URL(string: "\(baseURL)/api/auth/refresh")!
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = try JSONEncoder().encode(["refreshToken": currentRefreshToken])

        let maxRetries = 3
        var lastError: Error?

        for attempt in 0..<maxRetries {
            do {
                let (data, response) = try await URLSession.shared.data(for: request)

                guard let httpResponse = response as? HTTPURLResponse else {
                    throw AuthError.invalidResponse
                }

                if httpResponse.statusCode == 401 {
                    clearCredentials()
                    throw AuthError.refreshTokenExpired
                }

                guard (200...299).contains(httpResponse.statusCode) else {
                    throw AuthError.loginFailed(statusCode: httpResponse.statusCode)
                }

                let tokenResponse = try JSONDecoder.apiDecoder.decode(TokenResponse.self, from: data)
                saveKeychain(key: accessTokenKey, value: tokenResponse.accessToken)
                saveKeychain(key: refreshTokenKey, value: tokenResponse.refreshToken)

                lock.lock()
                _isAuthenticated = true
                lock.unlock()

                AppLogger.auth.info("Token refreshed successfully")
                return
            } catch let error as AuthError {
                throw error
            } catch {
                lastError = error
                if attempt < maxRetries - 1 {
                    let delay = UInt64(pow(2.0, Double(attempt))) * 1_000_000_000
                    try await Task.sleep(nanoseconds: delay)
                }
            }
        }

        clearCredentials()
        throw lastError.map { _ in AuthError.refreshFailed } ?? AuthError.refreshFailed
    }

    public func clearCredentials() {
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: keychainService,
        ]
        SecItemDelete(query as CFDictionary)

        lock.lock()
        _isAuthenticated = false
        lock.unlock()
    }

    // MARK: - SSO Login

    /// Exchange an SSO token for access and refresh tokens.
    /// - Parameter ssoToken: The SSO token received from the platform auth flow.
    /// - Throws: `AuthError` on failure.
    public func login(ssoToken: String) async throws {
        let url = URL(string: "\(baseURL)/api/auth/sso/exchange")!
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = try JSONEncoder().encode(["token": ssoToken])

        let (data, response) = try await URLSession.shared.data(for: request)

        guard let httpResponse = response as? HTTPURLResponse,
              (200...299).contains(httpResponse.statusCode) else {
            let statusCode = (response as? HTTPURLResponse)?.statusCode ?? 0
            throw AuthError.loginFailed(statusCode: statusCode)
        }

        let tokenResponse = try JSONDecoder.apiDecoder.decode(TokenResponse.self, from: data)
        saveKeychain(key: accessTokenKey, value: tokenResponse.accessToken)
        saveKeychain(key: refreshTokenKey, value: tokenResponse.refreshToken)

        guard isTokenValid(tokenResponse.accessToken) else {
            clearCredentials()
            throw AuthError.invalidToken
        }

        lock.lock()
        _isAuthenticated = true
        lock.unlock()

        AppLogger.auth.info("SSO login completed successfully")
    }

    /// Log out by clearing all stored credentials.
    public func logout() {
        clearCredentials()
        AppLogger.auth.info("User logged out")
    }

    // MARK: - JWT Validation

    private func isTokenValid(_ token: String) -> Bool {
        let segments = token.split(separator: ".")
        guard segments.count >= 2 else { return false }

        var base64 = String(segments[1])
        let remainder = base64.count % 4
        if remainder > 0 {
            base64 += String(repeating: "=", count: 4 - remainder)
        }

        guard let data = Data(base64Encoded: base64),
              let claims = try? JSONDecoder.apiDecoder.decode(JWTClaims.self, from: data) else {
            return false
        }

        // Check expiration if present
        if let exp = claims.exp {
            return Date().timeIntervalSince1970 < exp
        }

        return true
    }

    // MARK: - Keychain Operations

    private func saveKeychain(key: String, value: String) {
        let data = Data(value.utf8)
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: keychainService,
            kSecAttrAccount as String: key,
        ]

        SecItemDelete(query as CFDictionary)

        var addQuery = query
        addQuery[kSecValueData as String] = data
        addQuery[kSecAttrAccessible as String] = kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly

        SecItemAdd(addQuery as CFDictionary, nil)
    }

    private func readKeychain(key: String) -> String? {
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: keychainService,
            kSecAttrAccount as String: key,
            kSecReturnData as String: true,
            kSecMatchLimit as String: kSecMatchLimitOne,
        ]

        var result: AnyObject?
        let status = SecItemCopyMatching(query as CFDictionary, &result)

        guard status == errSecSuccess, let data = result as? Data else {
            return nil
        }

        return String(data: data, encoding: .utf8)
    }
}

// MARK: - Supporting Types

struct TokenResponse: Codable {
    let accessToken: String
    let refreshToken: String
}

struct JWTClaims: Codable {
    let sub: String
    let username: String?
    let displayName: String?
    let avatarURL: String?
    let provider: String?
    let exp: TimeInterval?
}