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

import Foundation

import Amplify
import AWSPluginsCore

protocol StorageMultipartUploadClient {
    func integrate(session: StorageMultipartUploadSession)

    func createMultipartUpload() throws
    func uploadPart(partNumber: PartNumber, multipartUpload: StorageMultipartUpload, subTask: StorageTransferTask) throws
    func completeMultipartUpload(uploadId: UploadID) throws
    func abortMultipartUpload(uploadId: UploadID, error: Error?) throws

    func cancelUploadTasks(taskIdentifiers: [TaskIdentifier], done: @escaping () -> Void)
}

extension StorageMultipartUploadClient {
    func abortMultipartUpload(uploadId: UploadID) throws {
        try abortMultipartUpload(uploadId: uploadId, error: nil)
    }
}

// Note: This may  be helpful in switching between Objective-C and Swift SDKs.
protocol StorageCreateMultipartUploadResponse {
    var uploadId: UploadID? { get }
}

class DefaultStorageMultipartUploadClient: StorageMultipartUploadClient {
    weak var serviceProxy: StorageServiceProxy?

    let fileSystem: FileSystem
    let uploadFile: UploadFile
    let bucket: String
    let key: String
    let contentType: String?
    let requestHeaders: RequestHeaders?
    weak var session: StorageMultipartUploadSession?
    let metadata: [String: String]?

    init(
        serviceProxy: StorageServiceProxy,
        fileSystem: FileSystem = .default,
        bucket: String,
        key: String,
        uploadFile: UploadFile,
        contentType: String? = nil,
        requestHeaders: RequestHeaders? = nil,
        metadata: [String: String]? = nil
    ) {
        self.serviceProxy = serviceProxy
        self.fileSystem = fileSystem
        self.bucket = bucket
        self.key = key
        self.uploadFile = uploadFile
        self.contentType = contentType
        self.requestHeaders = requestHeaders
        self.metadata = metadata
    }

    func integrate(session: StorageMultipartUploadSession) {
        self.session = session
    }

    // https://docs.aws.amazon.com/AmazonS3/latest/API/API_CreateMultipartUpload.html
    func createMultipartUpload() throws {
        guard let serviceProxy,
            let session else { fatalError() }

        // The AWS S3 SDK handles the request so there will be not taskIdentifier
        session.handle(multipartUploadEvent: .creating)

        // User-defined metadata needs to provided
        // when initiating the MPU.
        // --
        // https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html#mpu-process
        // > Multipart upload initiation
        // "If you want to provide any metadata describing the object
        // being uploaded, you must provide it in the request to initiate
        // multipart upload."
        let request = CreateMultipartUploadRequest(
            bucket: bucket,
            key: key,
            metadata: metadata
        )
        serviceProxy.awsS3.createMultipartUpload(request) { [weak self] result in
            guard let self else { return }
            switch result {
            case .success(let response):
                serviceProxy.register(multipartUploadSession: session)
                session.handle(multipartUploadEvent: .created(uploadFile: uploadFile, uploadId: response.uploadId))
            case .failure(let error):
                session.fail(error: error)
            }
        }
    }

    // https://docs.aws.amazon.com/AmazonS3/latest/API/API_UploadPart.html
    func uploadPart(partNumber: PartNumber, multipartUpload: StorageMultipartUpload, subTask: StorageTransferTask) throws {
        guard let serviceProxy else { fatalError("Service Proxy is required") }

        guard let uploadId = multipartUpload.uploadId,
              let uploadFile = multipartUpload.uploadFile,
              let partSize = multipartUpload.partSize,
              let part = multipartUpload.part(for: partNumber) else {
                  fatalError("Part number is required")
              }

        let startUploadPart: (URL, URL) async -> Void = { [weak self] partialFileURL, preSignedURL in
            guard let self else { return }
            var request = URLRequest(url: preSignedURL)
            request.cachePolicy = .reloadIgnoringLocalCacheData
            request.httpMethod = "PUT"
            request.networkServiceType = .responsiveData

            // Because this request represents an
            // [UploadPart](https://docs.aws.amazon.com/AmazonS3/latest/API/API_UploadPart.html)
            // operation, the "MultiPart/UploadPart" suffix is added to the
            // user agent.
            let userAgent = await serviceProxy.userAgent.appending(" MultiPart/UploadPart")
            request.setValue(userAgent, forHTTPHeaderField: "User-Agent")

            self.serviceProxy?.urlRequestDelegate?.willSend(request: request)
            let uploadTask = serviceProxy.urlSession.uploadTask(with: request, fromFile: partialFileURL)
            self.serviceProxy?.urlRequestDelegate?.didSend(request: request)

            subTask.sessionTask = uploadTask
            subTask.uploadPart = multipartUpload.part(for: partNumber)

            // register task so it can be found in delegate methods
            self.serviceProxy?.register(task: subTask)

            // tell the session the upload part has started
            session?.handle(uploadPartEvent: .started(partNumber: partNumber, taskIdentifier: uploadTask.taskIdentifier))

            uploadTask.resume()
        }

        let partialFileResultHandler: (Result<URL, Error>) -> Void = { [weak self] result in
            guard let self else { return }
            Task {
                do {
                    let partialFileURL = try result.get()
                    defer {
                        // Clean up the temporary file after we're done with it
                        self.fileSystem.removeFileIfExists(fileURL: partialFileURL)
                    }

                    let operation = AWSS3SigningOperation.uploadPart(partNumber: partNumber, uploadId: uploadId)
                    let preSignedURL = try await serviceProxy.preSignedURLBuilder.getPreSignedURL(
                        key: self.key,
                        signingOperation: operation,
                        // user-controlled metadata should *not* be provided
                        // with each upload part.
                        metadata: nil,
                        accelerate: nil,
                        expires: nil
                    )
                    await startUploadPart(partialFileURL, preSignedURL)
                } catch {
                    self.session?.fail(error: error)
                }
            }
        }

        let offset = partSize.offset(for: partNumber)
        fileSystem.createPartialFile(fileURL: uploadFile.fileURL, offset: offset, length: part.bytes, completionHandler: partialFileResultHandler)
    }

    // https://docs.aws.amazon.com/AmazonS3/latest/API/API_CompleteMultipartUpload.html
    func completeMultipartUpload(uploadId: UploadID) {
        guard let serviceProxy,
            let session else { fatalError() }

        let completedParts = session.completedParts ?? []
        let parts = AWSS3MultipartUploadRequestCompletedParts(completedParts: completedParts)
        let request = AWSS3CompleteMultipartUploadRequest(bucket: bucket, key: key, uploadId: uploadId, parts: parts)

        serviceProxy.awsS3.completeMultipartUpload(request) { result in
            switch result {
            case .success:
                session.handle(multipartUploadEvent: .completed(uploadId: uploadId))
            case .failure(let error):
                session.fail(error: error)
            }
            serviceProxy.unregister(multipartUploadSession: session)
        }
    }

    // https://docs.aws.amazon.com/AmazonS3/latest/API/API_AbortMultipartUpload.html
    func abortMultipartUpload(uploadId: UploadID, error: Error? = nil) {
        guard let serviceProxy,
            let session else { fatalError() }

        serviceProxy.awsS3.abortMultipartUpload(.init(bucket: bucket, key: key, uploadId: uploadId)) { result in
            switch result {
            case .success:
                session.handle(multipartUploadEvent: .aborted(uploadId: uploadId, error: error))
            case .failure(let error):
                session.fail(error: error)
            }
            serviceProxy.unregister(multipartUploadSession: session)
        }
    }

    func cancelUploadTasks(taskIdentifiers: [TaskIdentifier], done: @escaping () -> Void) {
        guard let service = serviceProxy else { return }
        service.unregister(taskIdentifiers: taskIdentifiers)
        service.urlSession.getActiveTasks { tasks in
            for task in tasks {
                if taskIdentifiers.contains(task.taskIdentifier) {
                    task.cancel()
                }
            }
            done()
        }
    }

    // MARK: - Private -

    // Note: the headers were previously filtered in the SDK
    func filter(requestHeaders: RequestHeaders) -> RequestHeaders {
        let disallowedHeaders: Set<String> = ["x-amz-acl", "x-amz-tagging", "x-amz-storage-class", "x-amz-server-side-encryption"]
        let shouldExcludeKey: (String) -> Bool = {
            $0.hasPrefix("x-amz-meta") ||
            $0.hasPrefix("x-amz-grant") ||
                disallowedHeaders.contains($0)
        }
        let result = requestHeaders.filter { key, _ in
            !shouldExcludeKey(key)
        }
        return result
    }

}
