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 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)
                )
                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)
        if let fileSize, fileSize > 524_288 {
            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
    }

    private func respondWithInMemoryData(fileURL: URL, requestURL: URL, taskID: ObjectIdentifier, urlSchemeTask: any WKURLSchemeTask) {
        do {
            let data = try Data(contentsOf: fileURL)
            guard isTaskActive(taskID) else { return }
            let response = URLResponse(
                url: requestURL,
                mimeType: Self.mimeType(for: fileURL.pathExtension),
                expectedContentLength: data.count,
                textEncodingName: Self.textEncodingName(for: fileURL.pathExtension)
            )
            urlSchemeTask.didReceive(response)
            urlSchemeTask.didReceive(data)
            urlSchemeTask.didFinish()
            Self.recordInMemoryResponse()
            RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, fileURL: fileURL, event: "finished")
        } catch {
            guard isTaskActive(taskID) else { return }
            Self.recordFailure()
            RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, fileURL: fileURL, event: "failed", error: error)
            urlSchemeTask.didFailWithError(error)
        }
        clearTask(taskID)
    }

    private func respondWithStreaming(fileURL: URL, requestURL: URL, taskID: ObjectIdentifier, urlSchemeTask: any WKURLSchemeTask, fileSize: UInt64) {
        let response = URLResponse(
            url: requestURL,
            mimeType: Self.mimeType(for: fileURL.pathExtension),
            expectedContentLength: Int(fileSize),
            textEncodingName: Self.textEncodingName(for: fileURL.pathExtension)
        )
        urlSchemeTask.didReceive(response)

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

        let chunkSize = 65_536
        defer {
            fileHandle.closeFile()
            clearTask(taskID)
        }
        while true {
            guard isTaskActive(taskID) else { return }
            let data = fileHandle.readData(ofLength: chunkSize)
            if data.isEmpty { break }
            urlSchemeTask.didReceive(data)
        }
        urlSchemeTask.didFinish()
        Self.recordStreamedResponse()
        RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, fileURL: fileURL, event: "finished-streaming")
    }

    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
        }
    }
}