ReadViewSDK/Sources/RDReaderView/EPUBCore/RDEPUBResourceURLSchemeHandler.swift

import Foundation
import WebKit

public final class RDEPUBResourceURLSchemeHandler: NSObject, WKURLSchemeHandler {
    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] = [:]

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

    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")
                urlSchemeTask.didFailWithError(NSError(domain: NSURLErrorDomain, code: NSURLErrorFileDoesNotExist))
            }
            clearTask(taskID)
            return
        }

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

        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()
            RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, fileURL: fileURL, event: "finished")
        } catch {
            guard isTaskActive(taskID) else { return }
            RDEPUBWebViewDebug.logSchemeTask("ResourceScheme", requestURL: requestURL, fileURL: fileURL, event: "failed", error: error)
            urlSchemeTask.didFailWithError(error)
        }

        clearTask(taskID)
    }

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