import Foundation

actor UploadManager {
    var progressContinuation: AsyncThrowingStream<Double, Error>.Continuation?
    var uploadTask: URLSessionUploadTask?
    var isCancelled = false
    var hasFinished = false

    func setProgressContinuation(_ continuation: AsyncThrowingStream<Double, Error>.Continuation?) {
        progressContinuation = continuation
    }

    func setUploadTask(_ task: URLSessionUploadTask?) {
        uploadTask = task
    }

    func yieldProgress(_ progress: Double) {
        guard !isCancelled else { return }
        progressContinuation?.yield(progress)
    }

    func finishProgress(with error: Error? = nil) {
        guard !hasFinished else { return }
        hasFinished = true
        if !isCancelled {
            if let error = error {
                progressContinuation?.finish(throwing: error)
            } else {
                progressContinuation?.finish()
            }
        }
    }

    func cancel() {
        isCancelled = true
        uploadTask?.cancel()
        if !hasFinished {
            hasFinished = true
            progressContinuation?.finish()
        }
    }

    func getIsCancelled() -> Bool {
        return isCancelled
    }
}

class SessionDelegate: NSObject, URLSessionDelegate, URLSessionTaskDelegate, URLSessionDataDelegate {
    let manager = UploadManager()
    let networkAccessIsLimited: (() -> Void)?
    weak var session: URLSession?

    init(networkAccessIsLimited: (() -> Void)?) {
        self.networkAccessIsLimited = networkAccessIsLimited
        super.init()
    }

    lazy var progress = AsyncThrowingStream<Double, Error> { [weak self] continuation in
        Task {
            await self?.manager.setProgressContinuation(continuation)
        }
    }

    func urlSession(_: URLSession, task _: URLSessionTask, didSendBodyData _: Int64, totalBytesSent: Int64, totalBytesExpectedToSend: Int64) {
        let progress = Double(totalBytesSent) / Double(totalBytesExpectedToSend)
        Task {
            await manager.yieldProgress(progress)
        }
    }

    func urlSession(_: URLSession, task _: URLSessionTask, didCompleteWithError error: Error?) {
        Task {
            if let error = error {
                if let urlError = error as? URLError {
                    if urlError.code == .cancelled {
                        await manager.finishProgress()
                        session?.finishTasksAndInvalidate()
                        return
                    }
                    if urlError.networkUnavailableReason == .cellular ||
                        urlError.networkUnavailableReason == .expensive
                    {
                        networkAccessIsLimited?()
                    }
                }
                log.telemetry.error(error, message: "S3 Upload failed")
                await manager.finishProgress(with: error)
            } else {
                await manager.finishProgress()
            }

            // Invalidate session to prevent memory leak
            session?.finishTasksAndInvalidate()
        }
    }

    func cancel() {
        Task {
            await manager.cancel()
        }
    }
}

public class S3UploadHandle {
    private weak var uploadTask: URLSessionUploadTask?

    init(uploadTask: URLSessionUploadTask?) {
        self.uploadTask = uploadTask
    }

    public func cancel() {
        uploadTask?.cancel()
    }
}

public struct S3UploadResult {
    public let progressStream: AsyncThrowingStream<Double, Error>
    public let handle: S3UploadHandle
}

func uploadToS3(
    localFileUrl: URL,
    uploadRequest: UploadRequest,
    limitUploadsOnLowNetwork: Bool,
    networkAccessIsLimited: (() -> Void)?
) throws -> S3UploadResult {
    let fileData = try Data(contentsOf: localFileUrl)
    return uploadDataToS3(
        fileData: fileData,
        filename: localFileUrl.lastPathComponent,
        uploadRequest: uploadRequest,
        limitUploadsOnLowNetwork: limitUploadsOnLowNetwork,
        networkAccessIsLimited: networkAccessIsLimited
    )
}

func uploadDataToS3(
    fileData: Data,
    filename: String,
    uploadRequest: UploadRequest,
    limitUploadsOnLowNetwork: Bool,
    networkAccessIsLimited: (() -> Void)?
) -> S3UploadResult {
    let delegate = SessionDelegate(networkAccessIsLimited: networkAccessIsLimited)
    var uploadHandle: S3UploadHandle?

    let progressStream = AsyncThrowingStream<Double, Error> { continuation in
        guard let uploadUrl = URL(string: uploadRequest.url) else {
            continuation.finish()
            return
        }

        var request = URLRequest(url: uploadUrl)
        request.httpMethod = "POST"
        request.setValue(uploadRequest.fields["Content-Type"], forHTTPHeaderField: "Content-Type")

        let boundary = UUID().uuidString
        Task {
            await delegate.manager.setProgressContinuation(continuation)
        }

        let config = URLSessionConfiguration.default
        if limitUploadsOnLowNetwork {
            config.allowsExpensiveNetworkAccess = false
            config.allowsCellularAccess = false
        }

        config.timeoutIntervalForRequest = Constants.timeoutIntervalForRequest
        config.timeoutIntervalForResource = Constants.timeoutIntervalForResource

        let session = URLSession(configuration: config, delegate: delegate, delegateQueue: nil)
        delegate.session = session

        request.setValue("multipart/form-data; boundary=\(boundary)", forHTTPHeaderField: "Content-Type")

        var data = Data()
        for (key, value) in uploadRequest.fields {
            if let dataValue = value.data(using: .utf8) {
                data.append("--\(boundary)\r\n".data(using: .utf8)!)
                data.append("Content-Disposition: form-data; name=\"\(key)\"\r\n\r\n".data(using: .utf8)!)
                data.append(dataValue)
                data.append("\r\n".data(using: .utf8)!)
            }
        }

        data.append("--\(boundary)\r\n".data(using: .utf8)!)
        data.append("Content-Disposition: form-data; name=\"file\"; filename=\"\(filename)\"\r\n".data(using: .utf8)!)
        data.append("Content-Type: \(uploadRequest.fields["Content-Type"] ?? "application/octet-stream")\r\n\r\n".data(using: .utf8)!)
        data.append(fileData)
        data.append("\r\n".data(using: .utf8)!)
        data.append("--\(boundary)--\r\n".data(using: .utf8)!)

        let task = session.uploadTask(with: request, from: data)

        Task {
            await delegate.manager.setUploadTask(task)
        }
        uploadHandle = S3UploadHandle(uploadTask: task)
        task.resume()

        continuation.onTermination = { _ in
            Task {
                await delegate.manager.cancel()
            }
            // Invalidate session when stream is terminated/cancelled
            session.invalidateAndCancel()
        }
    }

    return S3UploadResult(
        progressStream: progressStream,
        handle: uploadHandle ?? S3UploadHandle(uploadTask: nil)
    )
}


enum Constants {
    static let timeoutIntervalForRequest: TimeInterval = 60
    static let timeoutIntervalForResource: TimeInterval = 300
}