import Foundation

/// URLProtocol stub: tests enqueue responses, the client sees a real
/// URLSession. Recorded envelopes go through here.
///
/// Swift Testing runs suites in parallel, so state is keyed per box: each
/// `box()` gets its own queue and request log, matched up by a header the
/// box's session adds to every request.
nonisolated final class StubProtocol: URLProtocol, @unchecked Sendable {

    struct Stub: Sendable {
        let status: Int
        let headers: [String: String]
        let body: Data
        /// When set, this stub only answers requests whose URL contains
        /// it — needed when the code under test issues requests
        /// concurrently and arrival order is not deterministic.
        let match: String?

        init(status: Int, headers: [String: String] = [:], json: String, match: String? = nil) {
            self.status = status
            var headers = headers
            headers["Content-Type"] = headers["Content-Type"] ?? "application/json"
            self.headers = headers
            self.body = Data(json.utf8)
            self.match = match
        }

        init(status: Int, headers: [String: String] = [:], body: Data = Data(), match: String? = nil) {
            self.status = status
            self.headers = headers
            self.body = body
            self.match = match
        }
    }

    /// A request as the stub saw it, for asserting on method, URL, headers
    /// and body.
    struct Seen: Sendable {
        let url: URL
        let method: String
        let headers: [String: String]
        let body: Data
    }

    /// One test's private stub queue and request log.
    final class Box: Sendable {
        fileprivate let id = UUID().uuidString
        private let queue = Mutex<[Stub]>([])
        private let log = Mutex<[Seen]>([])

        func enqueue(_ stub: Stub) {
            queue.withLock { $0.append(stub) }
        }

        var seen: [Seen] {
            log.withLock { $0 }
        }

        /// The stubbed session for this box. The marker header is how the
        /// shared protocol class finds its way back here.
        func session() -> URLSession {
            let configuration = URLSessionConfiguration.ephemeral
            configuration.protocolClasses = [StubProtocol.self]
            configuration.httpAdditionalHeaders = [StubProtocol.marker: id]
            return URLSession(configuration: configuration)
        }

        fileprivate func next(for url: URL) -> Stub? {
            queue.withLock { stubs in
                let index = stubs.firstIndex {
                    $0.match.map { url.absoluteString.contains($0) } ?? true
                }
                guard let index else { return nil }
                return stubs.remove(at: index)
            }
        }

        fileprivate func record(_ request: Seen) {
            log.withLock { $0.append(request) }
        }
    }

    private static let marker = "X-Stub-Box"
    private static let boxes = Mutex<[String: Box]>([:])

    static func box() -> Box {
        let box = Box()
        boxes.withLock { $0[box.id] = box }
        return box
    }

    // MARK: - URLProtocol

    override static func canInit(with _: URLRequest) -> Bool { true }

    override static func canonicalRequest(for request: URLRequest) -> URLRequest { request }

    override func startLoading() {
        guard let id = request.value(forHTTPHeaderField: Self.marker),
              let box = Self.boxes.withLock({ $0[id] }) else {
            client?.urlProtocol(self, didFailWithError: URLError(.resourceUnavailable))
            return
        }

        // httpBody is transformed into a stream by URLSession; read it back.
        var body = Data()
        if let stream = request.httpBodyStream {
            stream.open()
            defer { stream.close() }
            let size = 4096
            let buffer = UnsafeMutablePointer<UInt8>.allocate(capacity: size)
            defer { buffer.deallocate() }
            while stream.hasBytesAvailable {
                let read = stream.read(buffer, maxLength: size)
                guard read > 0 else { break }
                body.append(buffer, count: read)
            }
        }
        var headers = request.allHTTPHeaderFields ?? [:]
        headers.removeValue(forKey: Self.marker)
        box.record(Seen(
            url: request.url!,
            method: request.httpMethod ?? "GET",
            headers: headers,
            body: body
        ))

        guard let stub = box.next(for: request.url!) else {
            client?.urlProtocol(self, didFailWithError: URLError(.resourceUnavailable))
            return
        }
        let response = HTTPURLResponse(
            url: request.url!,
            statusCode: stub.status,
            httpVersion: "HTTP/1.1",
            headerFields: stub.headers
        )!
        client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
        if !stub.body.isEmpty {
            client?.urlProtocol(self, didLoad: stub.body)
        }
        client?.urlProtocolDidFinishLoading(self)
    }

    override func stopLoading() {}
}

/// Minimal lock wrapper; os.OSAllocatedUnfairLock spelled test-side.
nonisolated final class Mutex<Value>: @unchecked Sendable {
    private var value: Value
    private let lock = NSLock()

    init(_ value: Value) {
        self.value = value
    }

    func withLock<R>(_ body: (inout Value) -> R) -> R {
        lock.lock()
        defer { lock.unlock() }
        return body(&value)
    }
}