//
//  APIRequestProxy.swift
//  MullvadVPN
//
//  Created by Jon Petersson on 2025-02-13.
//  Copyright © 2025 Mullvad VPN AB. All rights reserved.
//

import MullvadRustRuntime
import MullvadTypes

public protocol APIRequestProxyProtocol {
    func sendRequest(_ proxyRequest: ProxyAPIRequest, completion: @escaping @Sendable (ProxyAPIResponse) -> Void)
    func sendRequest(_ proxyRequest: ProxyAPIRequest) async -> ProxyAPIResponse
    func cancelRequest(identifier: UUID)
}

/// Network request proxy capable of passing serializable requests and responses over the given transport provider.
public final class APIRequestProxy: APIRequestProxyProtocol, @unchecked Sendable {
    /// Serial queue used for synchronizing access to class members.
    private let dispatchQueue: DispatchQueue

    private let transportProvider: APITransportProviderProtocol

    /// List of all proxied network requests bypassing VPN.
    private var proxiedRequests: [UUID: Cancellable] = [:]

    public init(
        dispatchQueue: DispatchQueue,
        transportProvider: APITransportProviderProtocol
    ) {
        self.dispatchQueue = dispatchQueue
        self.transportProvider = transportProvider
    }

    public func sendRequest(
        _ proxyRequest: ProxyAPIRequest,
        completion: @escaping @Sendable (ProxyAPIResponse) -> Void
    ) {
        dispatchQueue.async {
            guard let transport = self.transportProvider.makeTransport() else {
                // Cancel old task, if there's one scheduled.
                self.cancelRequest(identifier: proxyRequest.id)

                completion(ProxyAPIResponse(data: nil, error: nil))
                return
            }
            do {
                let cancellable = try transport.sendRequest(proxyRequest.request) { [weak self] response in
                    guard let self else { return }

                    // Use `dispatchQueue` to guarantee thread safe access to `proxiedRequests`
                    dispatchQueue.async {
                        _ = self.removeRequest(identifier: proxyRequest.id)
                        completion(response)
                    }
                }

                // Cancel old task, if there's one scheduled.
                let oldTask = self.addRequest(identifier: proxyRequest.id, task: cancellable)
                oldTask?.cancel()
            } catch {
                completion(
                    ProxyAPIResponse(
                        data: nil,
                        error: APIError(
                            statusCode: 0,
                            errorDescription: error.localizedDescription,
                            serverResponseCode: nil
                        )
                    )
                )
            }
        }
    }

    public func sendRequest(_ proxyRequest: ProxyAPIRequest) async -> ProxyAPIResponse {
        return await withCheckedContinuation { continuation in
            sendRequest(proxyRequest) { proxyResponse in
                continuation.resume(returning: proxyResponse)
            }
        }
    }

    public func cancelRequest(identifier: UUID) {
        dispatchQueue.async {
            let task = self.removeRequest(identifier: identifier)
            task?.cancel()
        }
    }

    private func addRequest(identifier: UUID, task: Cancellable) -> Cancellable? {
        dispatchPrecondition(condition: .onQueue(dispatchQueue))
        return proxiedRequests.updateValue(task, forKey: identifier)
    }

    private func removeRequest(identifier: UUID) -> Cancellable? {
        dispatchPrecondition(condition: .onQueue(dispatchQueue))
        return proxiedRequests.removeValue(forKey: identifier)
    }
}