ReadViewSDK/Sources/RDEpubReaderView/EPUBUI/ReaderController/RDEPUBBackgroundCoverageStore.swift

import Foundation

struct RDEPUBBackgroundCoverageSegment {

    let lowerSpineIndex: Int

    let upperSpineIndex: Int

    let pageMap: RDEPUBBookPageMap

    let resolvedSpineIndices: Set<Int>

    let generatedAt: CFAbsoluteTime

    let renderSignature: String

    let estimatedMemoryBytes: Int

    func contains(spineIndex: Int) -> Bool {
        spineIndex >= lowerSpineIndex && spineIndex <= upperSpineIndex
    }

    func distance(to spineIndex: Int) -> Int {
        if contains(spineIndex: spineIndex) { return 0 }
        return min(abs(spineIndex - lowerSpineIndex), abs(spineIndex - upperSpineIndex))
    }
}

struct RDEPUBBackgroundCoverageStorePolicy {

    let maxResidentSegments: Int

    let maxChaptersPerSegment: Int

    let memoryBudgetBytes: Int

    static let `default` = RDEPUBBackgroundCoverageStorePolicy(
        maxResidentSegments: 8,
        maxChaptersPerSegment: 256,
        memoryBudgetBytes: 8 * 1024 * 1024 
    )
}

final class RDEPUBBackgroundCoverageStore {

    private unowned let context: RDEPUBReaderContext

    private let policy: RDEPUBBackgroundCoverageStorePolicy

    private var segments: [RDEPUBBackgroundCoverageSegment] = []

    private var currentMemoryBytes: Int = 0

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

    init(context: RDEPUBReaderContext, policy: RDEPUBBackgroundCoverageStorePolicy = .default) {
        self.context = context
        self.policy = policy
    }

    func addSegment(_ segment: RDEPUBBackgroundCoverageSegment) {

        evictIfNeeded(forNewSegment: segment)

        var merged = false
        for (index, existing) in segments.enumerated() {
            if canMerge(existing, segment) {
                if let mergedSegment = mergeSegments(existing, segment) {
                    segments[index] = mergedSegment
                    currentMemoryBytes = currentMemoryBytes - existing.estimatedMemoryBytes + mergedSegment.estimatedMemoryBytes
                    merged = true
                    break
                }
            }
        }

        if !merged {
            segments.append(segment)
            currentMemoryBytes += segment.estimatedMemoryBytes
        }

        lastAccessTime[segment.lowerSpineIndex] = CFAbsoluteTimeGetCurrent()
    }

    func findSegment(containing spineIndex: Int) -> RDEPUBBackgroundCoverageSegment? {
        let segment = segments.first { $0.contains(spineIndex: spineIndex) }
        if let segment {
            lastAccessTime[segment.lowerSpineIndex] = CFAbsoluteTimeGetCurrent()
        }
        return segment
    }

    func findSegment(covering spineIndices: Set<Int>) -> RDEPUBBackgroundCoverageSegment? {
        let segment = segments.first { segment in
            spineIndices.allSatisfy { segment.contains(spineIndex: $0) }
        }
        if let segment {
            lastAccessTime[segment.lowerSpineIndex] = CFAbsoluteTimeGetCurrent()
        }
        return segment
    }

    func allSegments() -> [RDEPUBBackgroundCoverageSegment] {
        segments
    }

    func clearAll() {
        segments.removeAll()
        currentMemoryBytes = 0
        lastAccessTime.removeAll()
    }

    func clearColdSegments(
        activeWindowSpineIndices: Set<Int>,
        protectedSpineIndices: Set<Int>
    ) {
        let keepIndices = activeWindowSpineIndices.union(protectedSpineIndices)
        segments.removeAll { segment in
            let isCold = !segment.resolvedSpineIndices.contains(where: { keepIndices.contains($0) })
            if isCold {
                currentMemoryBytes -= segment.estimatedMemoryBytes
                lastAccessTime.removeValue(forKey: segment.lowerSpineIndex)
            }
            return isCold
        }
    }

    func handleMemoryWarning(
        activeWindowSpineIndices: Set<Int>,
        protectedSpineIndices: Set<Int>
    ) {

        clearColdSegments(
            activeWindowSpineIndices: activeWindowSpineIndices,
            protectedSpineIndices: protectedSpineIndices
        )

        if currentMemoryBytes > policy.memoryBudgetBytes {

            let sorted = segments.sorted { lhs, rhs in
                let lhsDistance = lhs.resolvedSpineIndices.map { idx in
                    activeWindowSpineIndices.map { abs(idx - $0) }.min() ?? Int.max
                }.min() ?? Int.max
                let rhsDistance = rhs.resolvedSpineIndices.map { idx in
                    activeWindowSpineIndices.map { abs(idx - $0) }.min() ?? Int.max
                }.min() ?? Int.max
                return lhsDistance > rhsDistance
            }

            for segment in sorted {
                if currentMemoryBytes <= policy.memoryBudgetBytes { break }
                currentMemoryBytes -= segment.estimatedMemoryBytes
                lastAccessTime.removeValue(forKey: segment.lowerSpineIndex)
                segments.removeAll { $0.lowerSpineIndex == segment.lowerSpineIndex }
            }
        }
    }

    private func evictIfNeeded(forNewSegment newSegment: RDEPUBBackgroundCoverageSegment) {

        while segments.count >= policy.maxResidentSegments {
            evictLeastRecentlyUsed()
        }

        while currentMemoryBytes + newSegment.estimatedMemoryBytes > policy.memoryBudgetBytes {
            evictLeastRecentlyUsed()
        }
    }

    private func evictLeastRecentlyUsed() {
        guard !segments.isEmpty else { return }

        var oldestTime = CFAbsoluteTimeGetCurrent()
        var oldestIndex = 0
        for (index, segment) in segments.enumerated() {
            let accessTime = lastAccessTime[segment.lowerSpineIndex] ?? 0
            if accessTime < oldestTime {
                oldestTime = accessTime
                oldestIndex = index
            }
        }

        let evicted = segments.remove(at: oldestIndex)
        currentMemoryBytes -= evicted.estimatedMemoryBytes
        lastAccessTime.removeValue(forKey: evicted.lowerSpineIndex)
    }

    private func canMerge(_ lhs: RDEPUBBackgroundCoverageSegment, _ rhs: RDEPUBBackgroundCoverageSegment) -> Bool {

        guard lhs.renderSignature == rhs.renderSignature else { return false }

        let overlap = lhs.upperSpineIndex >= rhs.lowerSpineIndex - 1 &&
                     rhs.upperSpineIndex >= lhs.lowerSpineIndex - 1
        return overlap
    }

    private func mergeSegments(_ lhs: RDEPUBBackgroundCoverageSegment, _ rhs: RDEPUBBackgroundCoverageSegment) -> RDEPUBBackgroundCoverageSegment? {
        let newLower = min(lhs.lowerSpineIndex, rhs.lowerSpineIndex)
        let newUpper = max(lhs.upperSpineIndex, rhs.upperSpineIndex)
        let newChapterCount = newUpper - newLower + 1

        if newChapterCount > policy.maxChaptersPerSegment {

            return nil
        }

        let newResolved = lhs.resolvedSpineIndices.union(rhs.resolvedSpineIndices)

        let newerSegment = lhs.generatedAt <= rhs.generatedAt ? rhs : lhs
        let olderSegment = lhs.generatedAt <= rhs.generatedAt ? lhs : rhs
        let newPageMap = mergePageMaps(olderSegment.pageMap, newerSegment.pageMap)

        return RDEPUBBackgroundCoverageSegment(
            lowerSpineIndex: newLower,
            upperSpineIndex: newUpper,
            pageMap: newPageMap,
            resolvedSpineIndices: newResolved,
            generatedAt: max(lhs.generatedAt, rhs.generatedAt),
            renderSignature: lhs.renderSignature,
            estimatedMemoryBytes: estimateMemoryBytes(pageMap: newPageMap, resolvedCount: newResolved.count)
        )
    }

    private func mergePageMaps(_ older: RDEPUBBookPageMap, _ newer: RDEPUBBookPageMap) -> RDEPUBBookPageMap {
        var builder = RDEPUBBookPageMap.Builder()
        var entriesBySpineIndex: [Int: RDEPUBBookPageMapEntry] = [:]

        for entry in older.entries {
            entriesBySpineIndex[entry.spineIndex] = entry
        }
        for entry in newer.entries {
            entriesBySpineIndex[entry.spineIndex] = entry
        }

        for spineIndex in entriesBySpineIndex.keys.sorted() {
            guard let entry = entriesBySpineIndex[spineIndex] else { continue }
            builder.add(
                spineIndex: entry.spineIndex,
                href: entry.href,
                title: entry.title,
                pageCount: entry.pageCount,
                fragmentOffsets: entry.fragmentOffsets
            )
        }

        return builder.build()
    }

    private func estimateMemoryBytes(pageMap: RDEPUBBookPageMap, resolvedCount: Int) -> Int {
        256 + pageMap.entries.count * 96 + resolvedCount * 16
    }
}