ReadViewSDK/Sources/RDReaderView/EPUBCore/RDEPUBResourceURLSchemeHandler.swift

import Foundation
import WebKit

public final class RDEPUBResourceURLSchemeHandler: NSObject, WKURLSchemeHandler {
    struct DebugMetrics {
        let streamedResponses: Int
        let inMemoryResponses: Int
        let failures: Int
    }

    public static let scheme = "ss-reader"

    public static let host = "book"

    private weak var parser: RDEPUBParser?

    private let fileManager = FileManager.default

    private let syncQueue = DispatchQueue(label: "com.ssreaderview.epub.scheme-handler")

    private let ioQueue = DispatchQueue(label: "com.ssreaderview.epub.scheme-handler.io", qos: .utility)

    private var activeTasks: [ObjectIdentifier: Bool] = [:]
    private static let debugMetricsQueue = DispatchQueue(label: "com.ssreaderview.epub.scheme-handler.metrics")
    private static var streamedResponseCount = 0
    private static var inMemoryResponseCount = 0
    private static var failureCount = 0

    public init(parser: RDEPUBParser) {
        self.parser = parser
        super.init()
    }

    static func resetDebugMetrics() {
        debugMetricsQueue.sync {
            streamedResponseCount = 0
            inMemoryResponseCount = 0
            failureCount = 0
        }
    }

    static func debugMetricsSnapshot() -> DebugMetrics {
        debugMetricsQueue.sync {
            DebugMetrics(
                streamedResponses: streamedResponseCount,
                inMemoryResponses: inMemoryResponseCount,
                failures: failureCount
            )
        }
    }

    public func webView(_ webView: WKWebView, start urlSchemeTask: any WKURLSchemeTask) {
        let taskID = ObjectIdentifier(urlSchemeTask as AnyObject)
        syncQueue.sync {
            activeTasks[taskID] = true
        }

        let requestURL = urlSchemeTask.request.url
        RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, event: "start")

        guard let parser,
              let requestURL,
              let fileURL = parser.fileURL(forResourceURL: requestURL),
              fileManager.fileExists(atPath: fileURL.path) else {
            if let requestURL, Self.isOptionalResource(requestURL) {
                RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, event: "missing-optional-resource")
                let response = URLResponse(
                    url: requestURL,
                    mimeType: Self.mimeType(for: requestURL.pathExtension),
                    expectedContentLength: 0,
                    textEncodingName: Self.textEncodingName(for: requestURL.pathExtension)
                )
                // H-08 fix: These calls happen synchronously in webView(_:start:) which
                // runs on the same queue WebKit calls start on, so no thread violation here.
                urlSchemeTask.didReceive(response)
                urlSchemeTask.didReceive(Data())
                urlSchemeTask.didFinish()
            } else {
                RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, event: "missing-file")
                Self.recordFailure()
                urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorFileDoesNotExist))
            }
            clearTask(taskID)
            return
        }

        RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, fileURL: fileURL, event: "resolved")

        let fileSize = resourceSize(for: fileURL)
        // 配置了加密资源 provider 时不能走流式分支：密文必须整体解密后再响应，
        // 统一改走内存分支（内部会先询问 provider，明文文件仍按磁盘直读）
        if let fileSize, fileSize > 524_288, parser.resourceDataProvider == nil {
            respondWithStreaming(fileURL: fileURL, requestURL: requestURL, taskID: taskID, urlSchemeTask: urlSchemeTask, fileSize: fileSize)
        } else {
            respondWithInMemoryData(fileURL: fileURL, requestURL: requestURL, taskID: taskID, urlSchemeTask: urlSchemeTask)
        }
    }

    public func webView(_ webView: WKWebView, stop urlSchemeTask: any WKURLSchemeTask) {
        RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: urlSchemeTask.request.url, event: "stop")
        clearTask(ObjectIdentifier(urlSchemeTask as AnyObject))
    }

    private func clearTask(_ id: ObjectIdentifier) {
        syncQueue.sync {
            activeTasks[id] = nil
        }
    }

    private func isTaskActive(_ id: ObjectIdentifier) -> Bool {
        syncQueue.sync {
            activeTasks[id] == true
        }
    }

    private func resourceSize(for url: URL) -> UInt64? {
        guard let attrs = try? fileManager.attributesOfItem(atPath: url.path),
              let size = attrs[.size] as? UInt64 else {
            return nil
        }
        return size
    }

    /// H-08 fix: Dispatch all WKURLSchemeTask completion callbacks to the main queue,
    /// ensuring they run on the same serial queue that WebKit calls start() on.
    /// Also: cancelled tasks now call didFailWithError(NSURLErrorCancelled) instead
    /// of silently returning, satisfying the protocol requirement that every started
    /// task must receive a completion callback.
    private func respondWithInMemoryData(fileURL: URL, requestURL: URL, taskID: ObjectIdentifier, urlSchemeTask: any WKURLSchemeTask) {
        ioQueue.async { [weak self] in
            guard let self else {
                DispatchQueue.main.async {
                    urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorCancelled))
                }
                return
            }
            do {
                // 优先向加密资源 provider 索取解密数据，provider 放行的明文文件仍直读磁盘
                let data: Data
                if let provided = self.parser?.providedResourceData(at: fileURL) {
                    data = provided
                } else {
                    data = try Data(contentsOf: fileURL)
                }
                guard self.isTaskActive(taskID) else {
                    DispatchQueue.main.async {
                        urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorCancelled))
                    }
                    return
                }
                let response = URLResponse(
                    url: requestURL,
                    mimeType: Self.mimeType(for: fileURL.pathExtension),
                    expectedContentLength: data.count,
                    textEncodingName: Self.textEncodingName(for: fileURL.pathExtension)
                )
                self.dispatchTaskSuccessCallback(
                    taskID: taskID,
                    urlSchemeTask: urlSchemeTask
                ) {
                    urlSchemeTask.didReceive(response)
                    urlSchemeTask.didReceive(data)
                    urlSchemeTask.didFinish()
                }
                Self.recordInMemoryResponse()
                RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, fileURL: fileURL, event: "finished")
                return
            } catch {
                guard self.isTaskActive(taskID) else {
                    DispatchQueue.main.async {
                        urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorCancelled))
                    }
                    return
                }
                Self.recordFailure()
                RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, fileURL: fileURL, event: "failed", error: error)
                DispatchQueue.main.async {
                    urlSchemeTask.didFailWithError(error)
                }
            }
            self.clearTask(taskID)
        }
    }

    private func respondWithStreaming(fileURL: URL, requestURL: URL, taskID: ObjectIdentifier, urlSchemeTask: any WKURLSchemeTask, fileSize: UInt64) {
        ioQueue.async { [weak self] in
            guard let self else {
                DispatchQueue.main.async {
                    urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorCancelled))
                }
                return
            }
            let response = URLResponse(
                url: requestURL,
                mimeType: Self.mimeType(for: fileURL.pathExtension),
                expectedContentLength: Int(fileSize),
                textEncodingName: Self.textEncodingName(for: fileURL.pathExtension)
            )
            self.dispatchTaskCallbackIfActive(taskID: taskID) {
                urlSchemeTask.didReceive(response)
            }

            guard let fileHandle = try? FileHandle(forReadingFrom: fileURL) else {
                guard self.isTaskActive(taskID) else {
                    DispatchQueue.main.async {
                        urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorCancelled))
                    }
                    return
                }
                Self.recordFailure()
                DispatchQueue.main.async {
                    urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorFileDoesNotExist))
                }
                self.clearTask(taskID)
                return
            }

            let chunkSize = 65_536
            defer {
                fileHandle.closeFile()
            }
            while true {
                guard self.isTaskActive(taskID) else {
                    // H-08 fix: Task was cancelled. Send didFailWithError on main queue.
                    DispatchQueue.main.async {
                        urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorCancelled))
                    }
                    self.clearTask(taskID)
                    return
                }
                let data = fileHandle.readData(ofLength: chunkSize)
                if data.isEmpty { break }
                self.dispatchTaskCallbackIfActive(taskID: taskID) {
                    urlSchemeTask.didReceive(data)
                }
            }
            self.dispatchTaskSuccessCallback(
                taskID: taskID,
                urlSchemeTask: urlSchemeTask
            ) {
                urlSchemeTask.didFinish()
            }
            Self.recordStreamedResponse()
            RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, fileURL: fileURL, event: "finished-streaming")
        }
    }

    private func dispatchTaskCallbackIfActive(taskID: ObjectIdentifier, _ callback: @escaping () -> Void) {
        DispatchQueue.main.async { [weak self] in
            guard let self, self.isTaskActive(taskID) else { return }
            callback()
        }
    }

    private func dispatchTaskSuccessCallback(
        taskID: ObjectIdentifier,
        urlSchemeTask: any WKURLSchemeTask,
        _ callback: @escaping () -> Void
    ) {
        DispatchQueue.main.async { [weak self] in
            guard let self else {
                urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorCancelled))
                return
            }
            guard self.isTaskActive(taskID) else {
                urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorCancelled))
                self.clearTask(taskID)
                return
            }
            callback()
            self.clearTask(taskID)
        }
    }

    private static func recordStreamedResponse() {
        debugMetricsQueue.sync {
            streamedResponseCount += 1
        }
    }

    private static func recordInMemoryResponse() {
        debugMetricsQueue.sync {
            inMemoryResponseCount += 1
        }
    }

    private static func recordFailure() {
        debugMetricsQueue.sync {
            failureCount += 1
        }
    }

    private static func mimeType(for pathExtension: String) -> String {
        switch pathExtension.lowercased() {
        case "html", "htm":
            return "text/html"
        case "xhtml":
            return "application/xhtml+xml"
        case "css":
            return "text/css"
        case "js":
            return "application/javascript"
        case "xml", "opf", "ncx":
            return "application/xml"
        case "svg":
            return "image/svg+xml"
        case "jpg", "jpeg":
            return "image/jpeg"
        case "png":
            return "image/png"
        case "gif":
            return "image/gif"
        case "webp":
            return "image/webp"
        case "ttf":
            return "font/ttf"
        case "otf":
            return "font/otf"
        case "woff":
            return "font/woff"
        case "woff2":
            return "font/woff2"
        case "mp3":
            return "audio/mpeg"
        case "mp4":
            return "video/mp4"
        default:
            return "application/octet-stream"
        }
    }

    private static func textEncodingName(for pathExtension: String) -> String? {
        switch pathExtension.lowercased() {
        case "html", "htm", "xhtml", "css", "js", "xml", "opf", "ncx", "txt":
            return "utf-8"
        default:
            return nil
        }
    }

    private static func isOptionalResource(_ url: URL) -> Bool {
        switch url.pathExtension.lowercased() {
        case "ttf", "otf", "woff", "woff2", "css", "js", "jpg", "jpeg", "png", "gif", "webp", "svg":
            return true
        default:
            return false
        }
    }
}