RetryableGraphQLOperation.swift

//
// Copyright Amazon.com Inc. or its affiliates.
// All Rights Reserved.
//
// SPDX-License-Identifier: Apache-2.0
//

import Foundation

/// Convenience protocol to handle any kind of GraphQLOperation
public protocol AnyGraphQLOperation {
    associatedtype Success
    associatedtype Failure: Error
    typealias ResultListener = (Result<Success, Failure>) -> Void
    typealias ErrorListener = (Failure) -> Void
}

/// Abastraction for a retryable GraphQLOperation.
public protocol RetryableGraphQLOperationBehavior: Operation, DefaultLogger {
    associatedtype Payload: Decodable

    /// GraphQLOperation concrete type
    associatedtype OperationType: AnyGraphQLOperation

    typealias RequestFactory = () async -> GraphQLRequest<Payload>
    typealias OperationFactory = (GraphQLRequest<Payload>, @escaping OperationResultListener) -> OperationType
    typealias OperationResultListener = OperationType.ResultListener
    typealias OperationErrorListener = OperationType.ErrorListener

    /// Operation unique identifier
    var id: UUID { get }

    /// Number of attempts (min 1)
    var attempts: Int { get set }

    /// Underlying GraphQL operation instantiated by `operationFactory`
    var underlyingOperation: AtomicValue<OperationType?> { get set }

    /// Maximum number of allowed retries
    var maxRetries: Int { get }

    /// GraphQLRequest factory, invoked to create a new operation
    var requestFactory: RequestFactory { get }

    /// GraphQL operation factory, invoked with a newly created GraphQL request
    /// and a wrapped result listener.
    var operationFactory: OperationFactory { get }

    var resultListener: OperationResultListener { get }
    
    var errorListener: OperationErrorListener { get }

    init(requestFactory: @escaping RequestFactory,
         maxRetries: Int,
         errorListener: @escaping OperationErrorListener,
         resultListener: @escaping OperationResultListener,
         _ operationFactory: @escaping OperationFactory)

    func start(request: GraphQLRequest<Payload>)

    func shouldRetry(error: APIError?) -> Bool
}

extension RetryableGraphQLOperationBehavior {
    public static var log: Logger {
        Amplify.Logging.logger(forCategory: CategoryType.api.displayName, forNamespace: String(describing: self))
    }
    public var log: Logger {
        Self.log
    }
}

// MARK: RetryableGraphQLOperationBehavior + default implementation
extension RetryableGraphQLOperationBehavior {
    public func start(request: GraphQLRequest<Payload>) {
        attempts += 1
        log.debug("[\(id)] - Try [\(attempts)/\(maxRetries)]")
        let wrappedResultListener: OperationResultListener = { result in
            if case let .failure(error) = result {
                // Give an operation a chance to prepare itself for a retry after a failure
                self.errorListener(error)
            }
            
            if case let .failure(error) = result, self.shouldRetry(error: error as? APIError) {
                self.log.debug("\(error)")
                Task {
                    self.start(request: await self.requestFactory())
                }
                return
            }

            if case let .failure(error) = result {
                self.log.debug("\(error)")
                self.log.debug("[\(self.id)] - Failed")
            }

            if case .success = result {
                self.log.debug("[Operation \(self.id)] - Success")
            }
            self.resultListener(result)
        }
        underlyingOperation.set(operationFactory(request, wrappedResultListener))
    }
}

// MARK: - RetryableGraphQLOperation
public final class RetryableGraphQLOperation<Payload: Decodable>: Operation, RetryableGraphQLOperationBehavior {
    public typealias Payload = Payload
    public typealias OperationType = GraphQLOperation<Payload>

    public var id: UUID
    public var maxRetries: Int
    public var attempts: Int = 0
    public var requestFactory: RequestFactory
    public var underlyingOperation: AtomicValue<GraphQLOperation<Payload>?> = AtomicValue(initialValue: nil)
    public var errorListener: OperationErrorListener
    public var resultListener: OperationResultListener
    public var operationFactory: OperationFactory

    public init(requestFactory: @escaping RequestFactory,
                maxRetries: Int,
                errorListener: @escaping OperationErrorListener,
                resultListener: @escaping OperationResultListener,
                _ operationFactory: @escaping OperationFactory) {
        self.id = UUID()
        self.maxRetries = max(1, maxRetries)
        self.requestFactory = requestFactory
        self.operationFactory = operationFactory
        self.errorListener = errorListener
        self.resultListener = resultListener
    }

    public override func main() {
        Task {
            start(request: await requestFactory())
        }
    }

    override public func cancel() {
        self.underlyingOperation.get()?.cancel()
    }

    public func shouldRetry(error: APIError?) -> Bool {
        guard case let .operationError(_, _, underlyingError) = error,
              let authError = underlyingError as? AuthError else {
                  return false
              }

        switch authError {
        case .signedOut, .notAuthorized:
            return attempts < maxRetries
        default:
            return false
        }
    }
}

// MARK: - RetryableGraphQLSubscriptionOperation
public final class RetryableGraphQLSubscriptionOperation<Payload: Decodable>: Operation,
                                                                              RetryableGraphQLOperationBehavior {
    public typealias OperationType = GraphQLSubscriptionOperation<Payload>

    public typealias Payload = Payload

    public var id: UUID
    public var maxRetries: Int
    public var attempts: Int = 0
    public var underlyingOperation: AtomicValue<GraphQLSubscriptionOperation<Payload>?> = AtomicValue(initialValue: nil)
    public var requestFactory: RequestFactory
    public var errorListener: OperationErrorListener
    public var resultListener: OperationResultListener
    public var operationFactory: OperationFactory
    private var retriedRTFErrors: [RTFError: Bool] = [:]
    
    public init(requestFactory: @escaping RequestFactory,
                maxRetries: Int,
                errorListener: @escaping OperationErrorListener,
                resultListener: @escaping OperationResultListener,
                _ operationFactory: @escaping OperationFactory) {
        self.id = UUID()
        self.maxRetries = max(1, maxRetries)
        self.requestFactory = requestFactory
        self.operationFactory = operationFactory
        self.errorListener = errorListener
        self.resultListener = resultListener
    }
    public override func main() {
        Task {
            start(request: await requestFactory())
        }
    }

    public override func cancel() {
        self.underlyingOperation.get()?.cancel()
    }

    public func shouldRetry(error: APIError?) -> Bool {
        guard case let .operationError(_, _, underlyingError) = error else {
            return false
        }
        
        if let authError = underlyingError as? AuthError {
            switch authError {
            case .signedOut, .notAuthorized:
                return attempts < maxRetries
            default:
                return false
            }
        }

        if let rtfError = RTFError(description: error.debugDescription) {
            
            // Do not retry the same RTF error more than once
            guard retriedRTFErrors[rtfError] == nil else { return false }
            retriedRTFErrors[rtfError] = true
                
            // maxRetries represent the number of auth types to attempt.
            // (maxRetries is set to the number of auth types to attempt in multi-auth rules scenarios)
            // Increment by 1 to account for that as this is not a "change auth" retry attempt            
            maxRetries += 1
                
            return true
        }
        
        return false
    }
}

// MARK: GraphQLOperation - GraphQLSubscriptionOperation + AnyGraphQLOperation
extension GraphQLOperation: AnyGraphQLOperation {}
extension GraphQLSubscriptionOperation: AnyGraphQLOperation {}