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

import Foundation
import XCTest
@testable import AWSDataStorePlugin

class RequestRetryablePolicyTests: XCTestCase {
    var retryPolicy: RequestRetryablePolicy!
    let defaultTimeout = DispatchTimeInterval.seconds(60)
    override func setUp() {
        super.setUp()
        retryPolicy = RequestRetryablePolicy()
    }

    func testNoErrorNoHttpURLResponse() {
        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: nil,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssertFalse(retryAdvice.shouldRetry)
        XCTAssertEqual(retryAdvice.retryInterval, defaultTimeout)
    }

    func testNoErrorWithHttpURLResponseWithRetryAfterInHeader() {
        let headerFields = ["Retry-After": "42"]
        let url = URL(string: "http://www.amazon.com")!
        let httpURLResponse = HTTPURLResponse(
            url: url,
            statusCode: 429,
            httpVersion: "HTTP/1.1",
            headerFields: headerFields
        )!

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: nil,
            httpURLResponse: httpURLResponse,
            attemptNumber: 1
        )
        XCTAssert(retryAdvice.shouldRetry)
        assertSeconds(retryAdvice.retryInterval, seconds: 42)
    }

    func testNoErrorWithHttpURLResponseWithoutRetryAfterInHeader_attempt1() {
        let url = URL(string: "http://www.amazon.com")!
        let httpURLResponse = HTTPURLResponse(
            url: url,
            statusCode: 429,
            httpVersion: "HTTP/1.1",
            headerFields: nil
        )!

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: nil,
            httpURLResponse: httpURLResponse,
            attemptNumber: 1
        )
        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testNoErrorWithHttpURLResponseWithoutRetryAfterInHeader_attempt2() {
        let url = URL(string: "http://www.amazon.com")!
        let httpURLResponse = HTTPURLResponse(
            url: url,
            statusCode: 429,
            httpVersion: "HTTP/1.1",
            headerFields: nil
        )!

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: nil,
            httpURLResponse: httpURLResponse,
            attemptNumber: 2
        )
        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 400, lessThan: 500)
    }

    func testNoErrorWithHttpURLResponseBeyondMaxWaitTime() {
        let url = URL(string: "http://www.amazon.com")!
        let httpURLResponse = HTTPURLResponse(
            url: url,
            statusCode: 429,
            httpVersion: "HTTP/1.1",
            headerFields: nil
        )!

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: nil,
            httpURLResponse: httpURLResponse,
            attemptNumber: 12
        )
        XCTAssertFalse(retryAdvice.shouldRetry)
        XCTAssertEqual(retryAdvice.retryInterval, defaultTimeout)
    }

    func testNoErrorWithHttpURLResponseNotRetryable() {
        let url = URL(string: "http://www.amazon.com")!
        let httpURLResponse = HTTPURLResponse(
            url: url,
            statusCode: 204,
            httpVersion: "HTTP/1.1",
            headerFields: nil
        )!

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: nil,
            httpURLResponse: httpURLResponse,
            attemptNumber: 1
        )
        XCTAssertFalse(retryAdvice.shouldRetry)
        XCTAssertEqual(retryAdvice.retryInterval, defaultTimeout)
    }

    func testNotConnectedToInternetErrorCode() {
        let retryableErrorCode = URLError.init(.notConnectedToInternet)
        let attemptNumber = 1

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: attemptNumber
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testNotConnectedToInternetErrorCode_attempt2() {
        let retryableErrorCode = URLError.init(.notConnectedToInternet)
        let attemptNumber = 2

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: attemptNumber
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 400, lessThan: 500)
    }

    func testNotConnectedToInternetErrorCode_attempt3() {
        let retryableErrorCode = URLError.init(.notConnectedToInternet)
        let attemptNumber = 3

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: attemptNumber
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 800, lessThan: 900)
    }

    func testDNSLookupFailedError() {
        let retryableErrorCode = URLError.init(.dnsLookupFailed)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testCannotConnectToHostError() {
        let retryableErrorCode = URLError.init(.cannotConnectToHost)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testCannotFindHostError() {
        let retryableErrorCode = URLError.init(.cannotFindHost)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testTimedOutError() {
        let retryableErrorCode = URLError.init(.timedOut)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testCannotParseResponseError() {
        let retryableErrorCode = URLError.init(.cannotParseResponse)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testNetworkConnectionLostError() {
        let retryableErrorCode = URLError.init(.networkConnectionLost)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testUserAuthenticationRequiredError() {
        let retryableErrorCode = URLError.init(.userAuthenticationRequired)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testHTTPTooManyRedirectsError() {
        let nonRetryableErrorCode = URLError.init(.httpTooManyRedirects)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: nonRetryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssertFalse(retryAdvice.shouldRetry)
        XCTAssertEqual(retryAdvice.retryInterval, defaultTimeout)
    }

    func testSecureConnectionFailedError() {
        let retryableErrorCode = URLError.init(.secureConnectionFailed)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 1
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
    }

    func testMaxValueRetryDelay() {
        let retryableErrorCode = URLError.init(.timedOut)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 31
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 214_748_364_800, lessThan: 214_748_364_900)
    }

    func testBeyondMaxValueRetryDelay() {
        let retryableErrorCode = URLError.init(.timedOut)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 32
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 214_748_364_800, lessThan: 214_748_364_900)
    }

    func testOverflowCase() {
        let retryableErrorCode = URLError.init(.timedOut)

        let retryAdvice = retryPolicy.retryRequestAdvice(
            urlError: retryableErrorCode,
            httpURLResponse: nil,
            attemptNumber: 58
        )

        XCTAssert(retryAdvice.shouldRetry)
        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 214_748_364_800, lessThan: 214_748_364_900)
    }

    func assertMilliseconds(_ retryInterval: DispatchTimeInterval?, greaterThan: Int, lessThan: Int) {
        guard let retryInterval else {
            XCTFail("retryInterval is nil")
            return
        }

        switch retryInterval {
        case .milliseconds(let milliseconds):
            XCTAssertGreaterThanOrEqual(milliseconds, greaterThan)
            XCTAssertLessThanOrEqual(milliseconds, lessThan)
        default:
            XCTFail("Expected milliseconds, but received \(retryInterval)")
        }
    }

    func assertSeconds(_ retryInterval: DispatchTimeInterval?, seconds expectedSeconds: Int) {
        guard let retryInterval else {
            XCTFail("retryInterval is nil")
            return
        }

        switch retryInterval {
        case .seconds(let seconds):
            XCTAssertEqual(seconds, expectedSeconds)
        default:
            XCTFail("Expected seconds, but received \(retryInterval)")
        }
    }
}
