ReadViewSDK/Sources/RDReaderView/EPUBUI/ReaderController/ChapterRuntime/RDEPUBChapterWarmupOrchestrator.swift

import Foundation

final class RDEPUBChapterWarmupOrchestrator {

    private unowned let context: RDEPUBReaderContext

    private unowned let store: RDEPUBChapterRuntimeStore

    private unowned let loader: RDEPUBChapterLoader

    private unowned let presentationRuntime: RDEPUBPresentationRuntime

    private unowned let locationCoordinator: RDEPUBReaderLocationCoordinator

    private unowned let backgroundPriorityManager: RDEPUBBackgroundPriorityManager

    private unowned let jumpSessionManager: RDEPUBJumpSessionManager

    private let refreshVisibleContentPreservingLocation: () -> Void

    private let asyncLoadStateLock = NSLock()

    private var asynchronouslyPreparingSpineIndices: Set<Int> = []

    private var isExtendingPartialBookPageMap = false

    private let prepareRequestStateLock = NSLock()

    private var pendingPreparePageNumbers: Set<Int> = []

    private var recentPrepareTimestamps: [Int: CFAbsoluteTime] = [:]

    private let prepareRequestDebounceInterval: CFTimeInterval = 0.15

    init(
        context: RDEPUBReaderContext,
        store: RDEPUBChapterRuntimeStore,
        loader: RDEPUBChapterLoader,
        presentationRuntime: RDEPUBPresentationRuntime,
        locationCoordinator: RDEPUBReaderLocationCoordinator,
        backgroundPriorityManager: RDEPUBBackgroundPriorityManager,
        jumpSessionManager: RDEPUBJumpSessionManager,
        refreshVisibleContentPreservingLocation: @escaping () -> Void
    ) {
        self.context = context
        self.store = store
        self.loader = loader
        self.presentationRuntime = presentationRuntime
        self.locationCoordinator = locationCoordinator
        self.backgroundPriorityManager = backgroundPriorityManager
        self.jumpSessionManager = jumpSessionManager
        self.refreshVisibleContentPreservingLocation = refreshVisibleContentPreservingLocation
    }

    @discardableResult
    func prepareOnDemandChapter(
        forAbsolutePageNumber pageNumber: Int,
        allowSynchronousLoad: Bool = true,
        completion: ((Bool) -> Void)? = nil
    ) -> Bool {
        guard let bookPageMap = context.bookPageMap,
              let publication = context.publication else {
            return false
        }
        let absolutePageIndex = pageNumber - 1
        guard let spineIndex = bookPageMap.spineIndex(forAbsolutePage: absolutePageIndex) else {
            return false
        }
        let chapterReady = store.chapterData(for: spineIndex) != nil

        if let debouncedResult = debouncedPrepareResult(
            pageNumber: pageNumber,
            spineIndex: spineIndex,
            chapterReady: chapterReady,
            allowSynchronousLoad: allowSynchronousLoad
        ) {
            return debouncedResult
        }

        store.setCurrentChapter(
            spineIndex: spineIndex,
            totalSpineCount: publication.spine.count,
            windowRadius: context.configuration.chapterWindowRadius
        )
        presentationRuntime.navigationStateMachine.transition(to: .preparingChapter(spineIndex: spineIndex))

        if !chapterReady {
            guard allowSynchronousLoad else {
                scheduleAsynchronousChapterPreparation(
                    spineIndex: spineIndex,
                    triggerPageNumber: pageNumber,
                    completion: completion
                )
                return false
            }
            do {
                _ = try loader.loadChapterSynchronouslyForMigration(
                    spineIndex: spineIndex,
                    store: store
                )
            } catch {
                clearPendingPreparePageNumber(pageNumber)
                return false
            }
        }

        markPrepareResolved(pageNumber)
        presentationRuntime.navigationStateMachine.transition(to: .presentingWindow)
        completion?(true)
        maybePrefetchUpcomingChapters(aroundAbsolutePageNumber: pageNumber, in: bookPageMap)
        scheduleAdjacentChapterPrefetches(for: spineIndex, totalSpineCount: publication.spine.count)
        return true
    }

    func extendPartialBookPageMapIfNeeded(
        currentPageNumber: Int,
        minimumTrailingPages: Int = 2,
        batchChapterCount: Int = 3
    ) {
        guard let publication = context.publication,
              let currentMap = context.bookPageMap else {
            return
        }

        let buildableSpineIndices = buildableSpineIndices(in: publication)
        guard currentMap.totalChapters < buildableSpineIndices.count else {
            return
        }

        let isNearEnd = currentMap.totalPages - currentPageNumber <= minimumTrailingPages
        let isNearStart = currentPageNumber <= minimumTrailingPages

        var spineIndicesToAppend: [Int] = []
        if isNearEnd {
            let lastKnownSpineIndex = currentMap.entries.last?.spineIndex ?? -1
            spineIndicesToAppend = Array(buildableSpineIndices.filter { $0 > lastKnownSpineIndex }.prefix(batchChapterCount))
        } else if isNearStart {
            let firstKnownSpineIndex = currentMap.entries.first?.spineIndex ?? Int.max
            let prependCandidates = buildableSpineIndices.filter { $0 < firstKnownSpineIndex }
            spineIndicesToAppend = Array(prependCandidates.suffix(batchChapterCount))
        } else {
            return
        }

        guard !spineIndicesToAppend.isEmpty, beginPartialBookPageMapExtension() else {
            return
        }


        let currentLocation = locationCoordinator.currentVisibleLocation()
        let loadedChaptersLock = NSLock()
        var loadedChapters: [Int: RDEPUBRuntimeChapter] = [:]
        let group = DispatchGroup()

        for spineIndex in spineIndicesToAppend {
            group.enter()
            loader.loadChapter(
                spineIndex: spineIndex,
                store: store,
                priority: .prefetch
            ) { result in
                defer { group.leave() }
                switch result {
                case .success(let chapter):
                    loadedChaptersLock.lock()
                    loadedChapters[spineIndex] = chapter
                    loadedChaptersLock.unlock()
                case .failure:
                    break
                }
            }
        }

        group.notify(queue: .main) { [weak self] in
            guard let self else { return }
            defer { self.endPartialBookPageMapExtension() }
            // Use live position values rather than the stale captured values,
            // since the user may have turned several pages since the extension began.
            let livePageNumber = max(self.context.readerView?.currentPage ?? 0, 0) + 1
            let liveLocation = self.locationCoordinator.currentVisibleLocation()
            self.applyAsyncPartialBookPageMapExtension(
                currentPageNumber: livePageNumber,
                currentLocation: liveLocation,
                currentMap: currentMap,
                loadedChapters: loadedChapters
            )
        }
    }

    func prefetchForwardChaptersAfterInitialOpen(anchorSpineIndex: Int, totalSpineCount: Int) {
        guard context.publication != nil else { return }

        store.setCurrentChapter(
            spineIndex: anchorSpineIndex,
            totalSpineCount: totalSpineCount,
            windowRadius: context.configuration.chapterWindowRadius
        )

        let forwardTargets = store.windowSpineIndices.filter { $0 > anchorSpineIndex }
        guard !forwardTargets.isEmpty else { return }

        for spineIndex in forwardTargets {
            if store.chapterData(for: spineIndex) != nil {
                appendLoadedForwardChaptersToCurrentPageMapIfPossible()
                continue
            }
            guard shouldSchedulePrefetch(for: spineIndex) else { continue }

            store.addPrefetchTarget(spineIndex)
            loader.loadChapter(spineIndex: spineIndex, store: store, priority: .prefetch) { [weak self] result in
                guard let self, case .success = result else { return }
                self.appendLoadedForwardChaptersToCurrentPageMapIfPossible()
            }
        }
    }

    @discardableResult
    func ensureNavigationTargetAvailable(for location: RDEPUBLocation) -> Bool {
        guard context.bookPageMap != nil,
              let publication = context.publication,
              let targetSpineIndex = context.normalizedSpineIndex(for: location) else {
            return false
        }

        let currentSpineIndex = locationCoordinator.currentVisibleLocation()
            .flatMap { context.normalizedSpineIndex(for: $0) }
        let isDistantJump = if let current = currentSpineIndex {
            abs(current - targetSpineIndex) > context.configuration.jumpSessionPolicy.protectedNeighborRadius * 2
        } else {
            false
        }

        if context.bookPageMap?.entry(forSpineIndex: targetSpineIndex) != nil {
            if isDistantJump {
                jumpSessionManager.createSession(
                    anchorSpineIndex: targetSpineIndex,
                    reason: .tableOfContentsJump,
                    totalSpineCount: publication.spine.count
                )
            }
            return true
        }

        if context.pendingPageMapUpdates.contains(where: { update in
            update.pageMap.entry(forSpineIndex: targetSpineIndex) != nil
        }) {
            presentationRuntime.commitPendingPageMapUpdateIfNeeded()
            if context.bookPageMap?.entry(forSpineIndex: targetSpineIndex) != nil {
                if isDistantJump {
                    jumpSessionManager.createSession(
                        anchorSpineIndex: targetSpineIndex,
                        reason: .tableOfContentsJump,
                        totalSpineCount: publication.spine.count
                    )
                }
                return true
            }
        }

        let buildableIndices = buildableSpineIndices(in: publication)
        guard let anchorPosition = buildableIndices.firstIndex(of: targetSpineIndex) else {
            return false
        }

        let normalizedWindowSize = RDEPUBReaderConfiguration.normalizedChapterWindowSize(
            context.configuration.onDemandChapterWindowSize
        )
        let chapters = loadPartialWindowChapters(
            around: anchorPosition,
            in: buildableIndices,
            targetSpineIndex: targetSpineIndex,
            windowSize: normalizedWindowSize
        )
        guard !chapters.isEmpty else { return false }

        store.setCurrentChapter(
            spineIndex: targetSpineIndex,
            totalSpineCount: publication.spine.count,
            windowRadius: context.configuration.chapterWindowRadius
        )
        let partialMap = makePartialPageMap(from: chapters)
        context.bookPageMap = partialMap
        context.replaceActiveSnapshot(presentationRuntime.makeSnapshot(from: partialMap))
        context.readerView?.reloadData()

        if isDistantJump {
            jumpSessionManager.createSession(
                anchorSpineIndex: targetSpineIndex,
                reason: .tableOfContentsJump,
                totalSpineCount: publication.spine.count
            )
            backgroundPriorityManager.addWarmAnchor(spineIndex: targetSpineIndex)
        }

        return partialMap.entry(forSpineIndex: targetSpineIndex) != nil
    }

    func clear() {
        asyncLoadStateLock.lock()
        asynchronouslyPreparingSpineIndices.removeAll()
        isExtendingPartialBookPageMap = false
        asyncLoadStateLock.unlock()
        prepareRequestStateLock.lock()
        pendingPreparePageNumbers.removeAll()
        recentPrepareTimestamps.removeAll()
        prepareRequestStateLock.unlock()
    }

    private func applyAsyncPartialBookPageMapExtension(
        currentPageNumber: Int,
        currentLocation: RDEPUBLocation?,
        currentMap: RDEPUBBookPageMap,
        loadedChapters: [Int: RDEPUBRuntimeChapter]
    ) {
        let appendedEntries = loadedChapters.keys.sorted().compactMap { spineIndex -> RDEPUBBookPageMapEntry? in
            guard let chapter = loadedChapters[spineIndex] else { return nil }
            return RDEPUBBookPageMapEntry(
                spineIndex: chapter.spineIndex,
                href: chapter.href,
                title: chapter.title,
                pageCount: chapter.pages.count,
                absolutePageStart: 0,
                fragmentOffsets: chapter.chapterOffsetMap.fragmentOffsets
            )
        }
        guard !appendedEntries.isEmpty else {
            return
        }

        let combinedEntries = (currentMap.entries.map {
            RDEPUBBookPageMapEntry(
                spineIndex: $0.spineIndex,
                href: $0.href,
                title: $0.title,
                pageCount: $0.pageCount,
                absolutePageStart: 0,
                fragmentOffsets: $0.fragmentOffsets
            )
        } + appendedEntries).sorted { $0.spineIndex < $1.spineIndex }

        var absolutePageStart = 0
        let normalizedEntries = combinedEntries.map { entry -> RDEPUBBookPageMapEntry in
            let normalized = RDEPUBBookPageMapEntry(
                spineIndex: entry.spineIndex,
                href: entry.href,
                title: entry.title,
                pageCount: entry.pageCount,
                absolutePageStart: absolutePageStart,
                fragmentOffsets: entry.fragmentOffsets
            )
            absolutePageStart += entry.pageCount
            return normalized
        }

        let newMap = RDEPUBBookPageMap(entries: normalizedEntries)
        presentationRuntime.queueExtendedPartialPageMap(
            newMap,
            currentPageNumber: currentPageNumber,
            currentLocation: currentLocation
        )
    }

    private func scheduleAsynchronousChapterPreparation(
        spineIndex: Int,
        triggerPageNumber: Int,
        completion: ((Bool) -> Void)?
    ) {
        guard beginAsynchronousChapterPreparation(for: spineIndex) else {
            return
        }
        loader.loadChapter(
            spineIndex: spineIndex,
            store: store,
            priority: .preview
        ) { [weak self] result in
            guard let self else { return }
            self.endAsynchronousChapterPreparation(for: spineIndex)
            switch result {
            case .success:
                self.markPrepareResolved(triggerPageNumber)
                self.presentationRuntime.navigationStateMachine.transition(to: .presentingWindow)
                completion?(true)
                self.appendLoadedForwardChaptersToCurrentPageMapIfPossible()
                self.refreshVisibleContentIfNeeded(afterPreparing: spineIndex, triggerPageNumber: triggerPageNumber)
            case .failure(let error):
                self.clearPendingPreparePageNumber(triggerPageNumber)
                completion?(false)
            }
        }
    }

    private func scheduleAdjacentChapterPrefetches(for spineIndex: Int, totalSpineCount: Int) {
        store.setCurrentChapter(
            spineIndex: spineIndex,
            totalSpineCount: totalSpineCount,
            windowRadius: context.configuration.chapterWindowRadius
        )

        for evictable in store.evictableSpineIndices() {
            store.evict(spineIndex: evictable)
        }

        for adjacentSpineIndex in store.windowSpineIndices where adjacentSpineIndex != spineIndex {
            guard shouldSchedulePrefetch(for: adjacentSpineIndex) else { continue }
            store.addPrefetchTarget(adjacentSpineIndex)
            loader.loadChapter(
                spineIndex: adjacentSpineIndex,
                store: store,
                priority: .prefetch
            ) { [weak self] result in
                guard let self, case .success = result else { return }
                self.appendLoadedForwardChaptersToCurrentPageMapIfPossible()
            }
        }
    }

    private func maybePrefetchUpcomingChapters(
        aroundAbsolutePageNumber pageNumber: Int,
        in bookPageMap: RDEPUBBookPageMap,
        threshold: Int = 3,
        lookaheadChapterCount: Int = 2
    ) {
        guard let publication = context.publication else { return }
        let absolutePageIndex = pageNumber - 1
        guard let spineIndex = bookPageMap.spineIndex(forAbsolutePage: absolutePageIndex),
              let localPageIndex = bookPageMap.localPageIndex(forAbsolutePage: absolutePageIndex),
              let chapter = store.chapterData(for: spineIndex) else {
            return
        }

        let remainingPages = chapter.pages.count - localPageIndex - 1
        guard remainingPages <= threshold else { return }

        let buildableIndices = buildableSpineIndices(in: publication)
        guard let currentPosition = buildableIndices.firstIndex(of: spineIndex) else { return }

        let targets = buildableIndices.dropFirst(currentPosition + 1).prefix(lookaheadChapterCount)
        for targetSpineIndex in targets {
            guard shouldSchedulePrefetch(for: targetSpineIndex) else { continue }
            store.addPrefetchTarget(targetSpineIndex)
            loader.loadChapter(spineIndex: targetSpineIndex, store: store, priority: .prefetch) { [weak self] result in
                guard let self, case .success = result else { return }
                self.appendLoadedForwardChaptersToCurrentPageMapIfPossible()
            }
        }
    }

    private func appendLoadedForwardChaptersToCurrentPageMapIfPossible(
        minimumTrailingPages: Int = 2
    ) {
        guard let publication = context.publication,
              let currentMap = context.bookPageMap,
              let readerView = context.readerView,
              let lastKnownSpineIndex = currentMap.entries.last?.spineIndex else {
            return
        }

        let currentPageNumber = max(readerView.currentPage + 1, 1)
        let trailingPages = currentMap.totalPages - currentPageNumber
        guard trailingPages <= minimumTrailingPages else { return }

        let buildableIndices = buildableSpineIndices(in: publication)
        var appendedEntries: [RDEPUBBookPageMapEntry] = []
        var projectedTotalPages = currentMap.totalPages
        for spineIndex in buildableIndices where spineIndex > lastKnownSpineIndex {
            guard let chapter = store.chapterData(for: spineIndex) else { break }
            appendedEntries.append(
                RDEPUBBookPageMapEntry(
                    spineIndex: chapter.spineIndex,
                    href: chapter.href,
                    title: chapter.title,
                    pageCount: chapter.pages.count,
                    absolutePageStart: 0,
                    fragmentOffsets: chapter.chapterOffsetMap.fragmentOffsets
                )
            )
            projectedTotalPages += chapter.pages.count
            if projectedTotalPages - currentPageNumber > minimumTrailingPages {
                break
            }
        }

        guard !appendedEntries.isEmpty else { return }

        let existingEntries = currentMap.entries.map {
            RDEPUBBookPageMapEntry(
                spineIndex: $0.spineIndex,
                href: $0.href,
                title: $0.title,
                pageCount: $0.pageCount,
                absolutePageStart: 0,
                fragmentOffsets: $0.fragmentOffsets
            )
        }

        var absolutePageStart = 0
        let newEntries = (existingEntries + appendedEntries).map { entry -> RDEPUBBookPageMapEntry in
            let normalizedEntry = RDEPUBBookPageMapEntry(
                spineIndex: entry.spineIndex,
                href: entry.href,
                title: entry.title,
                pageCount: entry.pageCount,
                absolutePageStart: absolutePageStart,
                fragmentOffsets: entry.fragmentOffsets
            )
            absolutePageStart += entry.pageCount
            return normalizedEntry
        }

        let newMap = RDEPUBBookPageMap(entries: newEntries)
        guard newMap.totalPages > currentMap.totalPages else { return }


        presentationRuntime.queueForwardAppendedPageMap(newMap)
    }

    private func shouldSchedulePrefetch(for spineIndex: Int) -> Bool {
        guard store.chapterData(for: spineIndex) == nil else { return false }
        guard !store.hasPrefetchTarget(spineIndex) else { return false }
        guard !store.hasPendingChapterLoad(for: spineIndex) else { return false }
        return true
    }

    private func debouncedPrepareResult(
        pageNumber: Int,
        spineIndex: Int,
        chapterReady: Bool,
        allowSynchronousLoad: Bool
    ) -> Bool? {
        prepareRequestStateLock.lock()
        defer { prepareRequestStateLock.unlock() }

        let now = CFAbsoluteTimeGetCurrent()
        recentPrepareTimestamps = recentPrepareTimestamps.filter { now - $0.value <= prepareRequestDebounceInterval }

        if !allowSynchronousLoad && !chapterReady {
            let inserted = pendingPreparePageNumbers.insert(pageNumber).inserted
            if !inserted {
                return false
            }
            return nil
        }

        pendingPreparePageNumbers.remove(pageNumber)
        if let lastTimestamp = recentPrepareTimestamps[pageNumber],
           now - lastTimestamp <= prepareRequestDebounceInterval {
            return chapterReady
        }

        recentPrepareTimestamps[pageNumber] = now
        return nil
    }

    private func markPrepareResolved(_ pageNumber: Int) {
        prepareRequestStateLock.lock()
        pendingPreparePageNumbers.remove(pageNumber)
        recentPrepareTimestamps[pageNumber] = CFAbsoluteTimeGetCurrent()
        prepareRequestStateLock.unlock()
    }

    private func clearPendingPreparePageNumber(_ pageNumber: Int) {
        prepareRequestStateLock.lock()
        pendingPreparePageNumbers.remove(pageNumber)
        prepareRequestStateLock.unlock()
    }

    private func refreshVisibleContentIfNeeded(afterPreparing spineIndex: Int, triggerPageNumber: Int) {
        guard let readerView = context.readerView,
              let bookPageMap = context.bookPageMap else {
            return
        }
        if readerView.isPageCurlTransitioning {
            DispatchQueue.main.asyncAfter(deadline: .now() + 0.1) { [weak self] in
                self?.refreshVisibleContentIfNeeded(
                    afterPreparing: spineIndex,
                    triggerPageNumber: triggerPageNumber
                )
            }
            return
        }
        let visiblePageNumber = readerView.currentPage + 1
        if visiblePageNumber == triggerPageNumber {
            refreshVisibleContentPreservingLocation()
            return
        }
        guard visiblePageNumber > 0,
              let visibleSpineIndex = bookPageMap.spineIndex(forAbsolutePage: visiblePageNumber - 1),
              visibleSpineIndex == spineIndex else {
            return
        }
        refreshVisibleContentPreservingLocation()
    }

    // Intentional synchronous path: called from TOC distant jumps where
    // the user expects immediate navigation. The loading indicator is shown
    // by the caller. Do NOT convert to async without UX consideration.
    private func loadPartialWindowChapters(
        around anchorPosition: Int,
        in buildableSpineIndices: [Int],
        targetSpineIndex: Int,
        windowSize: Int
    ) -> [RDEPUBRuntimeChapter] {
        let lowerBound = max(anchorPosition - max(windowSize / 2, 0), 0)
        let upperBound = min(lowerBound + max(windowSize, 1), buildableSpineIndices.count)
        let startIndex = max(0, upperBound - max(windowSize, 1))
        let window = Array(buildableSpineIndices[startIndex..<upperBound])

        var chapters: [RDEPUBRuntimeChapter] = []
        for spineIndex in window {
            do {
                let chapter = try loader.loadChapterSynchronouslyForMigration(
                    spineIndex: spineIndex,
                    store: store
                )
                chapters.append(chapter)
            } catch {
                if spineIndex == targetSpineIndex {
                    return []
                }
            }
        }
        return chapters
    }

    private func makePartialPageMap(from chapters: [RDEPUBRuntimeChapter]) -> RDEPUBBookPageMap {
        var builder = RDEPUBBookPageMap.Builder()
        for chapter in chapters {
            builder.add(
                spineIndex: chapter.spineIndex,
                href: chapter.href,
                title: chapter.title,
                pageCount: chapter.pages.count,
                fragmentOffsets: chapter.chapterOffsetMap.fragmentOffsets
            )
        }
        return builder.build()
    }

    private func buildableSpineIndices(in publication: RDEPUBPublication) -> [Int] {
        publication.spine.indices.filter { index in
            let item = publication.spine[index]
            return item.linear && (item.mediaType.contains("html") || item.mediaType.contains("xhtml"))
        }
    }

    private func beginAsynchronousChapterPreparation(for spineIndex: Int) -> Bool {
        asyncLoadStateLock.lock()
        defer { asyncLoadStateLock.unlock() }
        return asynchronouslyPreparingSpineIndices.insert(spineIndex).inserted
    }

    private func endAsynchronousChapterPreparation(for spineIndex: Int) {
        asyncLoadStateLock.lock()
        asynchronouslyPreparingSpineIndices.remove(spineIndex)
        asyncLoadStateLock.unlock()
    }

    private func beginPartialBookPageMapExtension() -> Bool {
        asyncLoadStateLock.lock()
        defer { asyncLoadStateLock.unlock() }
        guard !isExtendingPartialBookPageMap else { return false }
        isExtendingPartialBookPageMap = true
        return true
    }

    private func endPartialBookPageMapExtension() {
        asyncLoadStateLock.lock()
        isExtendingPartialBookPageMap = false
        asyncLoadStateLock.unlock()
    }
}