ReadViewSDK/Sources/RDReaderView/EPUBUI/ReaderController/RDEPUBTextStatisticsCoordinator.swift

import UIKit

struct RDEPUBChapterTextStatistics {

    let spineIndex: Int

    let href: String

    let characterCount: Int

    let absoluteCharacterStart: Int
}

struct RDEPUBBookTextStatistics {

    let chapters: [RDEPUBChapterTextStatistics]

    let totalCharacters: Int

    private let indexBySpine: [Int: Int]

    init(chapters: [RDEPUBChapterTextStatistics]) {
        self.chapters = chapters
        self.totalCharacters = chapters.reduce(0) { $0 + max($1.characterCount, 0) }
        self.indexBySpine = Dictionary(
            uniqueKeysWithValues: chapters.enumerated().map { ($0.element.spineIndex, $0.offset) }
        )
    }

    func progress(spineIndex: Int, chapterOffset: Int) -> Double? {
        guard totalCharacters > 0,
              let chapterIndex = indexBySpine[spineIndex] else {
            return nil
        }

        let chapter = chapters[chapterIndex]
        let safeChapterOffset = min(max(chapterOffset, 0), chapter.characterCount)
        let absoluteOffset = chapter.absoluteCharacterStart + safeChapterOffset
        return min(max(Double(absoluteOffset) / Double(totalCharacters), 0), 1)
    }

    func percent(spineIndex: Int, chapterOffset: Int) -> Double? {
        guard let progress = progress(spineIndex: spineIndex, chapterOffset: chapterOffset) else {
            return nil
        }
        return min(max(progress * 100, 0), 100)
    }
}

final class RDEPUBTextStatisticsWorker {

    private struct ChapterCharacterCount {

        let spineIndex: Int

        let href: String

        let characterCount: Int
    }

    func calculate(
        parser: RDEPUBParser,
        publication: RDEPUBPublication,
        bookIdentifier: String,
        renderSignature: String,
        style: RDEPUBTextRenderStyle,
        layoutConfig: RDEPUBTextLayoutConfig,
        rendererFactory: @escaping () -> RDEPUBTextRenderer,
        diskCache: RDEPUBTextStatisticsDiskCache?,
        maxConcurrency: Int,
        cancellationController: RDEPUBTextStatisticsCancellationController
    ) -> RDEPUBBookTextStatistics? {

        let buildableItems = publication.spine.enumerated().filter { _, item in
            isBuildableTextSpine(item)
        }

        let queue = OperationQueue()
        queue.name = "com.readview.epub.text-statistics"
        queue.qualityOfService = .utility
        queue.maxConcurrentOperationCount = max(1, maxConcurrency)
        cancellationController.attach(queue: queue)

        let lock = NSLock()
        var counts: [ChapterCharacterCount] = []

        for (spineIndex, item) in buildableItems {
            let operation = BlockOperation()
            operation.addExecutionBlock { [weak operation] in
                guard !cancellationController.isCancelled,
                      operation?.isCancelled != true else {
                    return
                }

                let characterCount = self.characterCount(
                    parser: parser,
                    publication: publication,
                    bookIdentifier: bookIdentifier,
                    renderSignature: renderSignature,
                    item: item,
                    spineIndex: spineIndex,
                    style: style,
                    layoutConfig: layoutConfig,
                    rendererFactory: rendererFactory,
                    diskCache: diskCache,
                    cancellationController: cancellationController
                )

                guard !cancellationController.isCancelled,
                      operation?.isCancelled != true else {
                    return
                }

                lock.lock()
                counts.append(
                    ChapterCharacterCount(
                        spineIndex: spineIndex,
                        href: item.href,
                        characterCount: characterCount
                    )
                )
                lock.unlock()
            }
            queue.addOperation(operation)
        }

        queue.waitUntilAllOperationsAreFinished()

        guard !cancellationController.isCancelled else {
            return nil
        }

        var entries: [RDEPUBChapterTextStatistics] = []
        var absoluteCharacterStart = 0
        for count in counts.sorted(by: { $0.spineIndex < $1.spineIndex }) {
            entries.append(
                RDEPUBChapterTextStatistics(
                    spineIndex: count.spineIndex,
                    href: count.href,
                    characterCount: count.characterCount,
                    absoluteCharacterStart: absoluteCharacterStart
                )
            )
            absoluteCharacterStart += max(count.characterCount, 0)
        }

        let statistics = RDEPUBBookTextStatistics(chapters: entries)
        return statistics.totalCharacters > 0 ? statistics : nil
    }

    private func characterCount(
        parser: RDEPUBParser,
        publication: RDEPUBPublication,
        bookIdentifier: String,
        renderSignature: String,
        item: RDEPUBSpineItem,
        spineIndex: Int,
        style: RDEPUBTextRenderStyle,
        layoutConfig: RDEPUBTextLayoutConfig,
        rendererFactory: () -> RDEPUBTextRenderer,
        diskCache: RDEPUBTextStatisticsDiskCache?,
        cancellationController: RDEPUBTextStatisticsCancellationController
    ) -> Int {
        autoreleasepool { () -> Int in
            guard let html = parser.htmlString(forRelativePath: item.href) else {
                return 0
            }

            let cacheKey = RDEPUBChapterCacheKey(
                bookID: bookIdentifier,
                spineIndex: spineIndex,
                renderSignature: renderSignature,
                chapterContentHash: html.rd_sha256Hex
            )
            if let cached = diskCache?.read(for: cacheKey) {
                return cached.characterCount
            }

            guard !cancellationController.isCancelled else {
                return 0
            }

            let renderer = rendererFactory()
            let baseURL = parser.fileURL(forRelativePath: item.href)?.deletingLastPathComponent()
            let request = RDEPUBTextRendererSupport.makeChapterRenderRequest(
                href: item.href,
                title: item.title,
                rawHTML: html,
                baseURL: baseURL,
                style: style,
                resourceResolver: publication.resourceResolver,
                contentLanguageCode: publication.metadata.language,
                pageSize: nil,
                layoutConfig: layoutConfig
            )

            do {
                let rendered = try renderer.renderChapter(request: request)
                let count = rendered.attributedString.length
                diskCache?.write(.init(characterCount: count), for: cacheKey)
                return count
            } catch {
                let fallback = RDEPUBTextRendererSupport.fallbackAttributedString(for: html, style: style)
                let count = fallback.length
                diskCache?.write(.init(characterCount: count), for: cacheKey)
                return count
            }
        }
    }

    private func isBuildableTextSpine(_ item: RDEPUBSpineItem) -> Bool {
        let mediaType = item.mediaType.lowercased()
        return item.linear && (mediaType.contains("html") || mediaType.contains("xhtml"))
    }
}

struct RDEPUBTextStatisticsCacheEntry: Codable {

    let schemaVersion: Int

    let characterCount: Int

    init(characterCount: Int, schemaVersion: Int = Self.currentSchemaVersion) {
        self.schemaVersion = schemaVersion
        self.characterCount = characterCount
    }

    static let currentSchemaVersion = 1
}

final class RDEPUBTextStatisticsDiskCache {

    private let cacheDirectory: URL

    private let fileManager = FileManager.default

    private let queue = DispatchQueue(label: "com.rdreader.textstatisticscache", qos: .utility)

    init(bookIdentifier: String?) {
        let cachesDirectory = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first
            ?? FileManager.default.temporaryDirectory
        let bookID = (bookIdentifier ?? "default").rd_sha256Hex
        self.cacheDirectory = cachesDirectory
            .appendingPathComponent("RDEPUBTextStatisticsCache", isDirectory: true)
            .appendingPathComponent(bookID, isDirectory: true)
        try? fileManager.createDirectory(at: cacheDirectory, withIntermediateDirectories: true)
    }

    func read(for key: RDEPUBChapterCacheKey) -> RDEPUBTextStatisticsCacheEntry? {
        let fileURL = fileURL(for: key)
        guard let data = try? Data(contentsOf: fileURL),
              let entry = try? JSONDecoder().decode(RDEPUBTextStatisticsCacheEntry.self, from: data),
              entry.schemaVersion == RDEPUBTextStatisticsCacheEntry.currentSchemaVersion else {
            return nil
        }
        return entry
    }

    func write(_ entry: RDEPUBTextStatisticsCacheEntry, for key: RDEPUBChapterCacheKey) {
        queue.async {
            self.writeImmediately(entry, for: key)
        }
    }

    func flushPendingWrites() {
        queue.sync { }
    }

    private func fileURL(for key: RDEPUBChapterCacheKey) -> URL {
        let bookPrefix = Self.cacheNamespacePrefix(for: key.bookID)
        let renderPrefix = Self.cacheNamespacePrefix(for: key.renderSignature)
        let rawKey = "\(key.bookID)_\(key.spineIndex)_\(key.renderSignature)_\(key.chapterContentHash)"
        let digest = rawKey.rd_sha256Hex
        return cacheDirectory.appendingPathComponent("\(bookPrefix)__\(renderPrefix)__\(digest).json")
    }

    private func writeImmediately(_ entry: RDEPUBTextStatisticsCacheEntry, for key: RDEPUBChapterCacheKey) {
        let fileURL = fileURL(for: key)
        let tmpURL = fileURL.appendingPathExtension("tmp")
        do {
            let data = try JSONEncoder().encode(entry)
            try data.write(to: tmpURL)
            if fileManager.fileExists(atPath: fileURL.path) {
                _ = try fileManager.replaceItemAt(fileURL, withItemAt: tmpURL)
            } else {
                try fileManager.moveItem(at: tmpURL, to: fileURL)
            }
        } catch {
            try? fileManager.removeItem(at: tmpURL)
        }
    }

    private static func cacheNamespacePrefix(for rawValue: String) -> String {
        rawValue.rd_sha256Hex.prefix(12).lowercased()
    }
}

final class RDEPUBTextStatisticsCancellationController {

    private let lock = NSLock()

    private var queue: OperationQueue?

    private var cancelled = false

    var isCancelled: Bool {
        lock.lock()
        defer { lock.unlock() }
        return cancelled
    }

    func attach(queue: OperationQueue) {
        lock.lock()
        if cancelled {
            queue.cancelAllOperations()
        } else {
            self.queue = queue
        }
        lock.unlock()
    }

    func cancel() {
        lock.lock()
        cancelled = true
        let queue = self.queue
        self.queue = nil
        lock.unlock()
        queue?.cancelAllOperations()
    }
}

final class RDEPUBTextStatisticsCoordinator {

    private weak var context: RDEPUBReaderContext?

    private let worker = RDEPUBTextStatisticsWorker()

    private let cancellationLock = NSLock()

    private var activeCancellationController: RDEPUBTextStatisticsCancellationController?

    init(context: RDEPUBReaderContext) {
        self.context = context
    }

    func start(parser: RDEPUBParser, publication: RDEPUBPublication) {
        guard let context else { return }

        guard publication.readingProfile == .textReflowable else {
            cancel()
            return
        }

        let token = UUID()
        cancelActiveWork()
        context.textStatisticsToken = token
        context.textStatistics = nil
        context.lastTextStatisticsWallClockMs = 0

        let snapshot = context.makeLayoutSnapshot()
        let dependencies = context.dependencies
        let renderingEngine = context.configuration.textRenderingEngine
        let workerCount = 1
        let bookIdentifier = context.currentBookIdentifier ?? ""
        let diskCache = RDEPUBTextStatisticsDiskCache(bookIdentifier: context.currentBookIdentifier)
        let cancellationController = RDEPUBTextStatisticsCancellationController()
        setActiveCancellationController(cancellationController)

        DispatchQueue.global(qos: .utility).async { [weak self, weak context] in
            guard let self else { return }
            let wallClockStart = CFAbsoluteTimeGetCurrent()
            let statistics = self.worker.calculate(
                parser: parser,
                publication: publication,
                bookIdentifier: bookIdentifier,
                renderSignature: snapshot.renderSignature,
                style: snapshot.style,
                layoutConfig: snapshot.layoutConfig,
                rendererFactory: { dependencies.makeTextRenderer(renderingEngine) },
                diskCache: diskCache,
                maxConcurrency: workerCount,
                cancellationController: cancellationController
            )
            diskCache.flushPendingWrites()
            let wallClockMs = Int((CFAbsoluteTimeGetCurrent() - wallClockStart) * 1000)

            DispatchQueue.main.async {
                guard let context,
                      context.textStatisticsToken == token,
                      !cancellationController.isCancelled else {
                    return
                }
                self.finishActiveCancellationController(cancellationController)
                context.textStatistics = statistics
                context.lastTextStatisticsWallClockMs = wallClockMs
                print("[TextStatistics] chapters=\(statistics?.chapters.count ?? 0) totalCharacters=\(statistics?.totalCharacters ?? 0) concurrency=\(workerCount) elapsed=\(wallClockMs)ms")
                context.readerView?.reloadPageCountOnly()
            }
        }
    }

    func cancel() {
        guard let context else { return }
        cancelActiveWork()
        context.textStatisticsToken = UUID()
        context.textStatistics = nil
        context.lastTextStatisticsWallClockMs = 0
    }

    private func cancelActiveWork() {
        cancellationLock.lock()
        let controller = activeCancellationController
        activeCancellationController = nil
        cancellationLock.unlock()
        controller?.cancel()
    }

    private func setActiveCancellationController(_ controller: RDEPUBTextStatisticsCancellationController) {
        cancellationLock.lock()
        activeCancellationController = controller
        cancellationLock.unlock()
    }

    private func finishActiveCancellationController(_ controller: RDEPUBTextStatisticsCancellationController) {
        cancellationLock.lock()
        if activeCancellationController === controller {
            activeCancellationController = nil
        }
        cancellationLock.unlock()
    }
}