ReadViewSDK/Sources/RDReaderView/EPUBTextRendering/RDEPUBChapterData.swift

import UIKit

public final class RDEPUBChapterData {

    public let chapter: RDEPUBTextChapter

    public let indexTable: RDEPUBTextIndexTable

    public init(chapter: RDEPUBTextChapter, indexTable: RDEPUBTextIndexTable) {
        self.chapter = chapter
        self.indexTable = indexTable
    }

    public var chapterIndex: Int { chapter.chapterIndex }

    public var spineIndex: Int { chapter.spineIndex }

    public var href: String { chapter.href }

    public var title: String { chapter.title }

    public var attributedContent: NSAttributedString { chapter.attributedContent }

    public var pages: [RDEPUBTextPage] { chapter.pages }

    public var pageCount: Int { chapter.pages.count }

    public var fragmentOffsets: [String: Int] { chapter.fragmentOffsets }

    public var chapterInfo: EPUBChapterInfo {
        EPUBChapterInfo(spineIndex: spineIndex, title: title, pageCount: pageCount)
    }

    public var absolutePageRange: ClosedRange<Int>? {
        guard let firstPage = pages.first, let lastPage = pages.last else { return nil }
        return firstPage.absolutePageIndex...lastPage.absolutePageIndex
    }

    public func page(containing absoluteOffset: Int) -> RDEPUBTextPage? {
        chapter.pages.first { NSLocationInRange(absoluteOffset, $0.contentRange) }
    }

    public func pageNumber(containing absoluteOffset: Int) -> Int? {
        page(containing: absoluteOffset)?.absolutePageIndex
    }

    public func page(atAbsolutePageIndex absolutePageIndex: Int) -> RDEPUBTextPage? {
        chapter.pages.first { $0.absolutePageIndex == absolutePageIndex }
    }

    public func page(atPageNumber pageNumber: Int) -> RDEPUBTextPage? {
        guard pageNumber > 0, pages.indices.contains(pageNumber - 1) else { return nil }
        return pages[pageNumber - 1]
    }

    public func anchor(forAbsoluteIndex index: Int) -> RDEPUBTextAnchor {
        indexTable.anchor(forAbsoluteIndex: index, in: chapter)
    }

    public func anchor(forGlobalIndex index: Int) -> RDEPUBTextAnchor? {
        indexTable.anchor(forGlobalIndex: index)
    }

    public func rangeAnchor(for absoluteRange: NSRange) -> RDEPUBTextRangeAnchor {
        let start = anchor(forAbsoluteIndex: absoluteRange.location)
        let end = anchor(forAbsoluteIndex: absoluteRange.location + absoluteRange.length)
        return RDEPUBTextRangeAnchor(start: start, end: end)
    }

    public func globalRange(for rangeAnchor: RDEPUBTextRangeAnchor) -> NSRange {
        indexTable.globalRange(for: rangeAnchor)
    }

    public func selection(from absoluteRange: NSRange, bookIdentifier: String?) -> RDEPUBSelection? {
        guard page(containing: absoluteRange.location) != nil else { return nil }
        let location = self.location(for: absoluteRange, bookIdentifier: bookIdentifier)
        let text = chapter.attributedContent.attributedSubstring(from: absoluteRange).string
        let rangeInfo = RDEPUBTextOffsetRangeInfo(
            href: chapter.href,
            start: absoluteRange.location,
            end: absoluteRange.location + absoluteRange.length
        ).jsonString()
        return RDEPUBSelection(
            bookIdentifier: bookIdentifier,
            location: location,
            text: text,
            rangeInfo: rangeInfo
        )
    }

    public func location(for absoluteRange: NSRange, bookIdentifier: String?) -> RDEPUBLocation {
        indexTable.location(
            for: rangeAnchor(for: absoluteRange),
            in: chapter,
            bookIdentifier: bookIdentifier
        )
    }

    public func location(forPage page: RDEPUBTextPage, bookIdentifier: String?) -> RDEPUBLocation {
        location(for: page.contentRange, bookIdentifier: bookIdentifier)
    }

    public func page(for location: RDEPUBLocation) -> RDEPUBTextPage? {
        guard let range = absoluteRange(for: location) else { return nil }
        return page(containing: range.location)
    }

    public func page(for searchMatch: RDEPUBSearchMatch) -> RDEPUBTextPage? {
        guard let range = absoluteRange(for: searchMatch) else { return nil }
        return page(containing: range.location)
    }

    public func absoluteRange(for location: RDEPUBLocation) -> NSRange? {
        if let cfiRange = RDEPUBCFICompatibility.parseRangeLossy(location.rangeCFI),
           let rangeAnchor = indexTable.rangeAnchor(for: cfiRange) {
            return indexTable.chapterRange(for: rangeAnchor)
        }

        if let cfi = RDEPUBCFICompatibility.parseLossy(location.cfi),
           let anchor = indexTable.anchor(for: cfi) {
            return NSRange(location: indexTable.chapterOffset(for: anchor), length: 1)
        }

        if let rangeAnchor = location.rangeAnchor {
            return indexTable.chapterRange(for: rangeAnchor)
        }

        if let fragment = location.fragment,
           let offset = fragmentOffsets[fragment] {
            return NSRange(location: offset, length: 1)
        }

        let lastOffset = max(attributedContent.length - 1, 0)
        let offset = min(lastOffset, max(0, Int(round(Double(lastOffset) * location.navigationProgression))))
        return NSRange(location: offset, length: 1)
    }

    public func absoluteRange(for highlight: RDEPUBHighlight) -> NSRange? {
        if let cfiRange = RDEPUBCFICompatibility.parseRangeLossy(highlight.location.rangeCFI),
           let rangeAnchor = indexTable.rangeAnchor(for: cfiRange) {
            return indexTable.chapterRange(for: rangeAnchor)
        }

        if let endpointRange = chapterRange(fromLocationCFIEndpoints: highlight.location) {
            return endpointRange
        }

        if let rangeAnchor = highlight.location.rangeAnchor {
            return indexTable.chapterRange(for: rangeAnchor)
        }
        return RDEPUBTextOffsetRangeInfo.decode(from: highlight.rangeInfo)?.nsRange
    }

    public func absoluteRange(for searchMatch: RDEPUBSearchMatch) -> NSRange? {
        if let cfiRange = RDEPUBCFICompatibility.parseRangeLossy(searchMatch.rangeCFI),
           let rangeAnchor = indexTable.rangeAnchor(for: cfiRange) {
            return indexTable.chapterRange(for: rangeAnchor)
        }
        if let cfi = RDEPUBCFICompatibility.parseLossy(searchMatch.cfi),
           let anchor = indexTable.anchor(for: cfi) {
            let location = indexTable.chapterOffset(for: anchor)
            return NSRange(
                location: location,
                length: recoveredSearchRangeLength(for: searchMatch, cfi: cfi, startOffset: location)
            )
        }
        if let location = searchMatch.rangeLocation {
            return NSRange(location: location, length: max(searchMatch.rangeLength, 1))
        }
        if let rangeAnchor = searchMatch.rangeAnchor {
            return indexTable.chapterRange(for: rangeAnchor)
        }
        return nil
    }

    public func globalRange(for location: RDEPUBLocation) -> NSRange? {
        if let cfiRange = RDEPUBCFICompatibility.parseRangeLossy(location.rangeCFI),
           let rangeAnchor = indexTable.rangeAnchor(for: cfiRange) {
            return indexTable.globalRange(for: rangeAnchor)
        }

        if let cfi = RDEPUBCFICompatibility.parseLossy(location.cfi),
           let anchor = indexTable.anchor(for: cfi) {
            return NSRange(location: indexTable.globalIndex(for: anchor), length: 1)
        }

        if let rangeAnchor = location.rangeAnchor {
            return indexTable.globalRange(for: rangeAnchor)
        }
        guard let chapterRange = absoluteRange(for: location),
              let chapterStart = indexTable.chapterStartOffset(forFileIndex: spineIndex) else {
            return nil
        }
        return NSRange(location: chapterStart + chapterRange.location, length: chapterRange.length)
    }

    public func globalRange(for highlight: RDEPUBHighlight) -> NSRange? {
        if let cfiRange = RDEPUBCFICompatibility.parseRangeLossy(highlight.location.rangeCFI),
           let rangeAnchor = indexTable.rangeAnchor(for: cfiRange) {
            return indexTable.globalRange(for: rangeAnchor)
        }

        if let endpointRange = globalRange(fromLocationCFIEndpoints: highlight.location) {
            return endpointRange
        }

        if let rangeAnchor = highlight.location.rangeAnchor {
            return indexTable.globalRange(for: rangeAnchor)
        }
        guard let chapterRange = absoluteRange(for: highlight),
              let chapterStart = indexTable.chapterStartOffset(forFileIndex: spineIndex) else {
            return nil
        }
        return NSRange(location: chapterStart + chapterRange.location, length: chapterRange.length)
    }

    public func globalRange(for searchMatch: RDEPUBSearchMatch) -> NSRange? {
        if let cfiRange = RDEPUBCFICompatibility.parseRangeLossy(searchMatch.rangeCFI),
           let rangeAnchor = indexTable.rangeAnchor(for: cfiRange) {
            return indexTable.globalRange(for: rangeAnchor)
        }
        if let cfi = RDEPUBCFICompatibility.parseLossy(searchMatch.cfi),
           let anchor = indexTable.anchor(for: cfi) {
            let chapterLocation = indexTable.chapterOffset(for: anchor)
            return NSRange(
                location: indexTable.globalIndex(for: anchor),
                length: recoveredSearchRangeLength(for: searchMatch, cfi: cfi, startOffset: chapterLocation)
            )
        }
        if let rangeAnchor = searchMatch.rangeAnchor {
            return indexTable.globalRange(for: rangeAnchor)
        }
        guard let chapterRange = absoluteRange(for: searchMatch),
              let chapterStart = indexTable.chapterStartOffset(forFileIndex: spineIndex) else {
            return nil
        }
        return NSRange(location: chapterStart + chapterRange.location, length: chapterRange.length)
    }

    public func highlights(on page: RDEPUBTextPage, from allHighlights: [RDEPUBHighlight]) -> [RDEPUBHighlight] {
        let pageRange = absoluteOffsetRange(for: page)
        return allHighlights.filter { highlight in
            guard highlight.location.href == chapter.href else { return false }
            if let cfiRange = RDEPUBCFICompatibility.parseRangeLossy(highlight.location.rangeCFI),
               let rangeAnchor = indexTable.rangeAnchor(for: cfiRange) {
                return NSIntersectionRange(indexTable.chapterRange(for: rangeAnchor), page.contentRange).length > 0
            }
            if let anchor = highlight.location.rangeAnchor?.start {
                return pageRange.contains(absoluteOffset(for: anchor))
            }
            guard let range = RDEPUBTextOffsetRangeInfo.decode(from: highlight.rangeInfo)?.nsRange else {
                return false
            }
            return NSIntersectionRange(range, page.contentRange).length > 0
        }
    }

    public func searchMatches(on page: RDEPUBTextPage, from matches: [RDEPUBSearchMatch]) -> [RDEPUBSearchMatch] {
        return matches.filter { match in
            guard match.href == chapter.href else { return false }
            if let range = absoluteRange(for: match) {
                return NSIntersectionRange(range, page.contentRange).length > 0
            }
            return false
        }
    }

    public func searchResults(on page: RDEPUBTextPage, from matches: [RDEPUBSearchMatch]) -> [RDEPUBSearchMatch] {
        searchMatches(on: page, from: matches)
    }

    public func pageNumber(for location: RDEPUBLocation) -> Int? {
        guard let page = page(for: location) else { return nil }
        return page.absolutePageIndex + 1
    }

    public func pageNumber(for searchMatch: RDEPUBSearchMatch) -> Int? {
        guard let page = page(for: searchMatch) else { return nil }
        return page.absolutePageIndex + 1
    }

    public func contains(
        tableOfContentsItem item: EPUBTableOfContentsItem,
        normalizer: (String) -> String?
    ) -> Bool {
        guard let chapterHref = normalizer(href),
              let itemHref = normalizer(item.href.components(separatedBy: "#").first ?? item.href) else {
            return false
        }
        return chapterHref == itemHref
    }

    public func tableOfContentsItems(
        from items: [EPUBTableOfContentsItem],
        normalizer: (String) -> String?
    ) -> [EPUBTableOfContentsItem] {
        items.flatMap { item -> [EPUBTableOfContentsItem] in
            let descendants = tableOfContentsItems(from: item.children, normalizer: normalizer)
            return contains(tableOfContentsItem: item, normalizer: normalizer) ? [item] + descendants : descendants
        }
    }

    public func primaryTableOfContentsItem(
        from items: [EPUBTableOfContentsItem],
        normalizer: (String) -> String?
    ) -> EPUBTableOfContentsItem? {
        tableOfContentsItems(from: items, normalizer: normalizer).first
    }

    public func applyHighlights(
        to content: NSMutableAttributedString,
        page: RDEPUBTextPage,
        highlights: [RDEPUBHighlight]
    ) {
        for highlight in highlights {
            guard highlight.location.href == chapter.href else { continue }
            guard let range = absoluteRange(for: highlight) else { continue }
            let overlap = NSIntersectionRange(range, page.contentRange)
            guard overlap.length > 0 else { continue }
            let relativeRange = NSRange(location: overlap.location - page.pageStartOffset, length: overlap.length)
            guard relativeRange.location >= 0,
                  relativeRange.location + relativeRange.length <= content.length else { continue }

            switch highlight.style {
            case .highlight:
                content.addAttribute(kRDEPUBHighlightAttributeName, value: highlight.uiColor, range: relativeRange)
            case .underline:
                content.addAttribute(kRDEPUBUnderlineAttributeName, value: highlight.uiColor, range: relativeRange)
            }
        }
    }

    private func absoluteOffsetRange(for page: RDEPUBTextPage) -> Range<Int> {
        let lowerBound = page.pageStartOffset
        let upperBound = page.pageEndOffset + 1
        return lowerBound..<max(upperBound, lowerBound)
    }

    private func absoluteOffset(for anchor: RDEPUBTextAnchor) -> Int {
        indexTable.chapterOffset(for: anchor)
    }

    private func chapterRange(fromLocationCFIEndpoints location: RDEPUBLocation) -> NSRange? {
        guard let startCFI = RDEPUBCFICompatibility.parseLossy(location.cfi),
              let endCFI = RDEPUBCFICompatibility.parseLossy(location.lastCFI ?? location.cfi),
              let startAnchor = indexTable.anchor(for: startCFI),
              let endAnchor = indexTable.anchor(for: endCFI),
              startAnchor.fileIndex == endAnchor.fileIndex else {
            return nil
        }
        let start = indexTable.chapterOffset(for: startAnchor)
        let end = indexTable.chapterOffset(for: endAnchor)
        return NSRange(location: min(start, end), length: max(abs(end - start), 1))
    }

    private func globalRange(fromLocationCFIEndpoints location: RDEPUBLocation) -> NSRange? {
        guard let startCFI = RDEPUBCFICompatibility.parseLossy(location.cfi),
              let endCFI = RDEPUBCFICompatibility.parseLossy(location.lastCFI ?? location.cfi),
              let startAnchor = indexTable.anchor(for: startCFI),
              let endAnchor = indexTable.anchor(for: endCFI),
              startAnchor.fileIndex == endAnchor.fileIndex else {
            return nil
        }
        let start = indexTable.globalIndex(for: startAnchor)
        let end = indexTable.globalIndex(for: endAnchor)
        return NSRange(location: min(start, end), length: max(abs(end - start), 1))
    }

    private func recoveredSearchRangeLength(
        for searchMatch: RDEPUBSearchMatch,
        cfi: RDEPUBCFI,
        startOffset: Int
    ) -> Int {
        let exactLength = cfi.textAssertion?.exact?.utf16.count ?? 0
        let fallbackLength = max(searchMatch.rangeLength, 1)
        let candidateLength = exactLength > 0 ? exactLength : fallbackLength
        let remainingLength = max(attributedContent.length - startOffset, 1)
        return min(max(candidateLength, 1), remainingLength)
    }
}