ReadViewSDK/Sources/RDReaderView/EPUBTextRendering/RDEPUBTextIndexTable.swift

import Foundation

public struct RDEPUBRowColumnIndex: Codable, Equatable {

    public let row: Int

    public let startOffset: Int

    public let endOffset: Int

    public init(row: Int, startOffset: Int, endOffset: Int) {
        self.row = row
        self.startOffset = startOffset
        self.endOffset = endOffset
    }

    public func contains(_ offset: Int) -> Bool {
        offset >= startOffset && offset <= endOffset
    }
}

public struct RDEPUBTextIndexTable {

    public let chapterStartOffsets: [Int]

    public let chapterLengths: [Int]

    public let hrefToChapterIndex: [String: Int]

    public let hrefToFileIndex: [String: Int]

    public let fragmentOffsetsByHref: [String: [String: Int]]

    public let fileIndexToHref: [Int: String]

    public let fileRowColumnMap: [Int: [RDEPUBRowColumnIndex]]

    public let fileTextMap: [Int: String]

    public let fileCFIMap: [Int: RDEPUBCFIMap]

    public var totalCharacterCount: Int {
        guard let lastIndex = chapterStartOffsets.indices.last,
              chapterLengths.indices.contains(lastIndex) else {
            return 0
        }
        return chapterStartOffsets[lastIndex] + chapterLengths[lastIndex]
    }

    public init(chapters: [RDEPUBTextChapter]) {
        var offsets: [Int] = []
        var lengths: [Int] = []
        var hrefMap: [String: Int] = [:]
        var hrefToFileMap: [String: Int] = [:]
        var fragmentMap: [String: [String: Int]] = [:]
        var fileIndexMap: [Int: String] = [:]
        var rowColumnMap: [Int: [RDEPUBRowColumnIndex]] = [:]
        var textMap: [Int: String] = [:]
        var cfiMap: [Int: RDEPUBCFIMap] = [:]
        var running = 0

        for (index, chapter) in chapters.enumerated() {
            hrefMap[chapter.href] = index
            hrefToFileMap[chapter.href] = chapter.spineIndex
            offsets.append(running)
            lengths.append(chapter.attributedContent.length)
            running += chapter.attributedContent.length
            fragmentMap[chapter.href] = chapter.fragmentOffsets
            fileIndexMap[chapter.spineIndex] = chapter.href
            let chapterText = chapter.attributedContent.string
            rowColumnMap[chapter.spineIndex] = Self.makeRowColumnIndices(for: chapterText)
            textMap[chapter.spineIndex] = chapterText
            if let chapterCFIMap = chapter.cfiMap {
                cfiMap[chapter.spineIndex] = chapterCFIMap
            }
        }

        self.chapterStartOffsets = offsets
        self.chapterLengths = lengths
        self.hrefToChapterIndex = hrefMap
        self.hrefToFileIndex = hrefToFileMap
        self.fragmentOffsetsByHref = fragmentMap
        self.fileIndexToHref = fileIndexMap
        self.fileRowColumnMap = rowColumnMap
        self.fileTextMap = textMap
        self.fileCFIMap = cfiMap
    }

    public func anchor(forAbsoluteIndex index: Int, in chapter: RDEPUBTextChapter) -> RDEPUBTextAnchor {
        let normalizedIndex = clampedOffset(index, in: chapter)
        let fragmentID = nearestFragmentID(beforeOrAt: normalizedIndex, in: chapter)
        let row = row(forAbsoluteIndex: normalizedIndex, inFileIndex: chapter.spineIndex)
        let column = column(forAbsoluteIndex: normalizedIndex, inFileIndex: chapter.spineIndex)
        return RDEPUBTextAnchor(
            fileIndex: chapter.spineIndex,
            row: row,
            column: column,
            chapterOffset: normalizedIndex,
            fragmentID: fragmentID
        )
    }

    public func anchor(for location: RDEPUBLocation) -> RDEPUBTextAnchor? {
        if let cfiRange = RDEPUBCFICompatibility.parseRangeLossy(location.rangeCFI),
           let anchor = anchor(for: cfiRange.start) {
            return anchor
        }

        if let anchor = location.rangeAnchor?.start {
            return anchor
        }

        if let cfi = RDEPUBCFICompatibility.parseLossy(location.cfi),
           let anchor = anchor(for: cfi) {
            return anchor
        }

        guard let chapterIndex = hrefToChapterIndex[location.href],
              let fileIndex = hrefToFileIndex[location.href] else { return nil }
        let fragments = fragmentOffsetsByHref[location.href] ?? [:]
        let chapterOffset: Int

        if let fragment = location.fragment, let fragmentOffset = fragments[fragment] {
            chapterOffset = fragmentOffset
        } else {
            let estimatedLength = max(chapterLengths.indices.contains(chapterIndex) ? chapterLengths[chapterIndex] : 0, 1)
            let lastOffset = max(estimatedLength - 1, 0)
            chapterOffset = min(lastOffset, max(0, Int(round(Double(lastOffset) * location.navigationProgression))))
        }

        return RDEPUBTextAnchor(
            fileIndex: fileIndex,
            row: row(forAbsoluteIndex: chapterOffset, inFileIndex: fileIndex),
            column: column(forAbsoluteIndex: chapterOffset, inFileIndex: fileIndex),
            chapterOffset: chapterOffset,
            fragmentID: location.fragment
        )
    }

    public func anchor(for cfi: RDEPUBCFI) -> RDEPUBTextAnchor? {
        let resolved = RDEPUBCFIResolver.resolve(cfi)
        let fileIndex: Int
        if let resolvedFileIndex = resolved.fileIndex {
            fileIndex = resolvedFileIndex
        } else if let href = resolved.href, let mappedFileIndex = hrefToFileIndex[href] {
            fileIndex = mappedFileIndex
        } else {
            return nil
        }

        guard let href = href(for: fileIndex) ?? resolved.href else { return nil }
        let chapterLength = max(chapterLength(forFileIndex: fileIndex) ?? 0, 1)
        let lastOffset = max(chapterLength - 1, 0)
        let recovery = RDEPUBCFIRecoveryEngine.recover(
            cfi: cfi,
            cfiMap: fileCFIMap[fileIndex],
            chapterText: fileTextMap[fileIndex],
            fragmentOffsets: fragmentOffsetsByHref[href] ?? [:],
            fallbackOffset: resolved.chapterOffset,
            lastOffset: lastOffset
        )
        let chapterOffset = recovery?.chapterOffset ?? min(max(resolved.chapterOffset ?? 0, 0), lastOffset)
        let fragmentID = resolved.fragmentID ?? nearestFragmentID(beforeOrAt: chapterOffset, inHref: href)

        return RDEPUBTextAnchor(
            fileIndex: fileIndex,
            row: row(forAbsoluteIndex: chapterOffset, inFileIndex: fileIndex),
            column: column(forAbsoluteIndex: chapterOffset, inFileIndex: fileIndex),
            chapterOffset: chapterOffset,
            fragmentID: fragmentID
        )
    }

    public func rangeAnchor(for cfiRange: RDEPUBCFIRange) -> RDEPUBTextRangeAnchor? {
        guard let start = anchor(for: cfiRange.start),
              let end = anchor(for: cfiRange.end),
              start.fileIndex == end.fileIndex else {
            return nil
        }
        return RDEPUBTextRangeAnchor(start: start, end: end)
    }

    public func cfi(for anchor: RDEPUBTextAnchor) -> RDEPUBCFI? {
        guard let href = href(for: anchor.fileIndex) else { return nil }
        let chapterOffset = chapterOffset(for: anchor)
        if let marker = preciseMarker(fileIndex: anchor.fileIndex, chapterOffset: chapterOffset) {
            let localOffset = max(chapterOffset - (marker.chapterOffset ?? 0), 0)
            return RDEPUBCFIGenerator.makeCFI(
                href: href,
                fileIndex: anchor.fileIndex,
                contentPath: marker.cfiPath,
                characterOffset: localOffset,
                textAssertion: textAssertion(fileIndex: anchor.fileIndex, offset: chapterOffset)
            )
        }
        return RDEPUBCFIGenerator.makeOffsetCFI(
            href: href,
            fileIndex: anchor.fileIndex,
            chapterOffset: chapterOffset,
            fragmentID: anchor.fragmentID,
            textAssertion: textAssertion(fileIndex: anchor.fileIndex, offset: chapterOffset)
        )
    }

    public func cfiRange(for rangeAnchor: RDEPUBTextRangeAnchor) -> RDEPUBCFIRange? {
        guard rangeAnchor.start.fileIndex == rangeAnchor.end.fileIndex,
              let href = href(for: rangeAnchor.start.fileIndex) else {
            return nil
        }
        let startOffset = chapterOffset(for: rangeAnchor.start)
        let endOffset = chapterOffset(for: rangeAnchor.end)
        let startCFI = cfi(for: rangeAnchor.start)
        let endCFI = cfi(for: rangeAnchor.end)
        if let startCFI, let endCFI {
            let parent = RDEPUBCFI(
                packagePath: startCFI.packagePath.commonPrefix(with: endCFI.packagePath),
                contentPath: startCFI.contentPath.commonPrefix(with: endCFI.contentPath)
            )
            return RDEPUBCFIRange(parent: parent, start: startCFI, end: endCFI)
        }
        return RDEPUBCFIGenerator.makeOffsetRangeCFI(
            href: href,
            fileIndex: rangeAnchor.start.fileIndex,
            startOffset: startOffset,
            endOffset: endOffset,
            fragmentID: rangeAnchor.start.fragmentID,
            startTextAssertion: textAssertion(
                fileIndex: rangeAnchor.start.fileIndex,
                offset: startOffset
            ),
            endTextAssertion: textAssertion(
                fileIndex: rangeAnchor.end.fileIndex,
                offset: endOffset
            )
        )
    }

    public func anchor(forGlobalIndex index: Int) -> RDEPUBTextAnchor? {
        guard let fileIndex = fileIndex(forCharacterPosition: index),
              let href = href(for: fileIndex) else {
            return nil
        }

        let chapterOffset = localOffsetInFile(at: fileIndex, forGlobalPosition: index) ?? 0
        let fragmentID = nearestFragmentID(beforeOrAt: chapterOffset, inHref: href)
        return RDEPUBTextAnchor(
            fileIndex: fileIndex,
            row: row(forAbsoluteIndex: chapterOffset, inFileIndex: fileIndex),
            column: column(forAbsoluteIndex: chapterOffset, inFileIndex: fileIndex),
            chapterOffset: chapterOffset,
            fragmentID: fragmentID
        )
    }

    public func pageNumber(for anchor: RDEPUBTextAnchor, in book: RDEPUBTextBook) -> Int? {
        guard let chapter = book.chapters.first(where: { $0.spineIndex == anchor.fileIndex }) else { return nil }
        let basePageIndex = chapter.pages.first?.absolutePageIndex ?? 0
        let resolvedOffset = chapterOffset(for: anchor)
        return chapter.pages.firstIndex { page in
            NSLocationInRange(resolvedOffset, page.contentRange)
        }.map { $0 + basePageIndex }
    }

    public func href(for fileIndex: Int) -> String? {
        fileIndexToHref[fileIndex]
    }

    public func chapterIndex(for href: String) -> Int? {
        hrefToChapterIndex[href]
    }

    public func chapterStartOffset(forFileIndex fileIndex: Int) -> Int? {
        guard let href = href(for: fileIndex),
              let chapterIndex = hrefToChapterIndex[href],
              chapterStartOffsets.indices.contains(chapterIndex) else {
            return nil
        }
        return chapterStartOffsets[chapterIndex]
    }

    public func chapterLength(forFileIndex fileIndex: Int) -> Int? {
        guard let href = href(for: fileIndex),
              let chapterIndex = hrefToChapterIndex[href],
              chapterLengths.indices.contains(chapterIndex) else {
            return nil
        }
        return chapterLengths[chapterIndex]
    }

    public func chapterOffset(for anchor: RDEPUBTextAnchor) -> Int {
        anchor.chapterOffset
    }

    public func globalIndex(for anchor: RDEPUBTextAnchor) -> Int {
        let chapterOffset = chapterOffset(for: anchor)
        let baseOffset = chapterStartOffset(forFileIndex: anchor.fileIndex) ?? 0
        let chapterLength = max(chapterLength(forFileIndex: anchor.fileIndex) ?? 0, 0)
        let lastOffset = max(chapterLength - 1, 0)
        return baseOffset + min(max(chapterOffset, 0), lastOffset)
    }

    public func absoluteIndex(for anchor: RDEPUBTextAnchor) -> Int {
        globalIndex(for: anchor)
    }

    public func chapterRange(for rangeAnchor: RDEPUBTextRangeAnchor) -> NSRange {
        let start = chapterOffset(for: rangeAnchor.start)
        let end = max(start, chapterOffset(for: rangeAnchor.end))
        return NSRange(location: start, length: max(end - start, 0))
    }

    public func globalRange(for rangeAnchor: RDEPUBTextRangeAnchor) -> NSRange {
        let start = globalIndex(for: rangeAnchor.start)
        let end = max(start, globalIndex(for: rangeAnchor.end))
        return NSRange(location: start, length: max(end - start, 0))
    }

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

    public func location(
        for anchor: RDEPUBTextAnchor,
        in chapter: RDEPUBTextChapter,
        bookIdentifier: String?
    ) -> RDEPUBLocation {
        let chapterOffset = self.chapterOffset(for: anchor)
        let totalLength = max(chapter.attributedContent.length - 1, 1)
        let progression = Double(min(max(chapterOffset, 0), totalLength)) / Double(totalLength)
        return RDEPUBLocation(
            bookIdentifier: bookIdentifier,
            href: chapter.href,
            progression: progression,
            lastProgression: progression,
            fragment: anchor.fragmentID,
            rangeAnchor: RDEPUBTextRangeAnchor(start: anchor, end: anchor),
            cfi: cfi(for: anchor)?.rawValue,
            lastCFI: cfi(for: anchor)?.rawValue
        )
    }

    public func location(
        for rangeAnchor: RDEPUBTextRangeAnchor,
        in chapter: RDEPUBTextChapter,
        bookIdentifier: String?
    ) -> RDEPUBLocation {
        let start = chapterOffset(for: rangeAnchor.start)
        let end = max(start, chapterOffset(for: rangeAnchor.end))
        let totalLength = max(chapter.attributedContent.length - 1, 1)
        let clampedStart = min(max(start, 0), totalLength)
        let clampedEnd = min(max(end, clampedStart), totalLength)
        return RDEPUBLocation(
            bookIdentifier: bookIdentifier,
            href: chapter.href,
            progression: Double(clampedStart) / Double(totalLength),
            lastProgression: Double(clampedEnd) / Double(totalLength),
            fragment: rangeAnchor.start.fragmentID,
            rangeAnchor: rangeAnchor,
            cfi: cfi(for: rangeAnchor.start)?.rawValue,
            lastCFI: cfi(for: rangeAnchor.end)?.rawValue,
            rangeCFI: cfiRange(for: rangeAnchor)?.rawValue
        )
    }

    public func row(forAbsoluteIndex index: Int, inFileIndex fileIndex: Int) -> Int {
        guard let rows = fileRowColumnMap[fileIndex], !rows.isEmpty else { return 0 }
        if let rowIndex = rows.first(where: { $0.contains(index) })?.row {
            return rowIndex
        }
        return rows.last?.row ?? 0
    }

    public func column(forAbsoluteIndex index: Int, inFileIndex fileIndex: Int) -> Int {
        guard let rows = fileRowColumnMap[fileIndex], !rows.isEmpty else { return 0 }
        if let rowEntry = rows.first(where: { $0.contains(index) }) {
            return max(index - rowEntry.startOffset, 0)
        }
        guard let lastRow = rows.last else { return 0 }
        return max(index - lastRow.startOffset, 0)
    }

    public func chapterOffset(fileIndex: Int, row: Int, column: Int) -> Int? {
        guard let rows = fileRowColumnMap[fileIndex], !rows.isEmpty else { return nil }
        let normalizedRow = min(max(row, 0), rows.count - 1)
        let rowEntry = rows[normalizedRow]
        let maxColumn = max(rowEntry.endOffset - rowEntry.startOffset, 0)
        return rowEntry.startOffset + min(max(column, 0), maxColumn)
    }

    public func absoluteIndex(fileIndex: Int, row: Int, column: Int) -> Int? {
        guard let chapterOffset = chapterOffset(fileIndex: fileIndex, row: row, column: column),
              let baseOffset = chapterStartOffset(forFileIndex: fileIndex) else {
            return nil
        }
        return baseOffset + chapterOffset
    }

    public func fileIndex(forCharacterPosition index: Int) -> Int? {
        guard totalCharacterCount > 0 else { return nil }
        let normalized = min(max(index, 0), max(totalCharacterCount - 1, 0))
        for (href, chapterIndex) in hrefToChapterIndex {
            guard chapterStartOffsets.indices.contains(chapterIndex),
                  chapterLengths.indices.contains(chapterIndex),
                  let fileIndex = hrefToFileIndex[href] else {
                continue
            }

            let start = chapterStartOffsets[chapterIndex]
            let length = chapterLengths[chapterIndex]
            let endExclusive = start + max(length, 1)
            if normalized >= start && normalized < endExclusive {
                return fileIndex
            }
        }
        return fileIndexToHref.keys.sorted().last
    }

    public func localOffsetInFile(at fileIndex: Int, forGlobalPosition index: Int) -> Int? {
        guard let start = chapterStartOffset(forFileIndex: fileIndex),
              let chapterLength = chapterLength(forFileIndex: fileIndex) else {
            return nil
        }
        let normalized = min(max(index, start), start + max(chapterLength - 1, 0))
        return normalized - start
    }

    private func clampedOffset(_ index: Int, in chapter: RDEPUBTextChapter) -> Int {
        let lastOffset = max(chapter.attributedContent.length - 1, 0)
        return min(max(index, 0), lastOffset)
    }

    private func nearestFragmentID(beforeOrAt offset: Int, in chapter: RDEPUBTextChapter) -> String? {
        var bestID: String?
        var bestOffset = -1

        for (id, fragOffset) in chapter.fragmentOffsets {
            if fragOffset <= offset && fragOffset > bestOffset {
                bestOffset = fragOffset
                bestID = id
            }
        }

        return bestID
    }

    private func nearestFragmentID(beforeOrAt offset: Int, inHref href: String) -> String? {
        guard let fragmentOffsets = fragmentOffsetsByHref[href] else {
            return nil
        }

        var bestID: String?
        var bestOffset = -1
        for (id, fragmentOffset) in fragmentOffsets where fragmentOffset <= offset && fragmentOffset > bestOffset {
            bestOffset = fragmentOffset
            bestID = id
        }
        return bestID
    }

    private func preciseMarker(fileIndex: Int, chapterOffset: Int) -> RDEPUBCFIMarker? {
        guard let markers = fileCFIMap[fileIndex]?.markers else { return nil }
        return markers
            .filter {
                guard let start = $0.chapterOffset,
                      let length = $0.textNodeLength,
                      length > 0 else {
                    return false
                }
                return chapterOffset >= start && chapterOffset < start + length
            }
            .max { lhs, rhs in
                let leftStart = lhs.chapterOffset ?? 0
                let rightStart = rhs.chapterOffset ?? 0
                if leftStart != rightStart {
                    return leftStart < rightStart
                }
                return (lhs.textNodeLength ?? 0) > (rhs.textNodeLength ?? 0)
            }
    }

    private func textAssertion(fileIndex: Int, offset: Int) -> RDEPUBCFITextAssertion? {
        guard let text = fileTextMap[fileIndex], !text.isEmpty else { return nil }
        let nsText = text as NSString
        let length = nsText.length
        guard length > 0 else { return nil }

        let location = min(max(offset, 0), max(length - 1, 0))
        let exactLength = min(16, max(length - location, 0))
        guard exactLength > 0 else { return nil }

        let prefixStart = max(location - 8, 0)
        let suffixStart = min(location + exactLength, length)
        return RDEPUBCFITextAssertion(
            prefix: prefixStart < location
                ? nsText.substring(with: NSRange(location: prefixStart, length: location - prefixStart))
                : nil,
            exact: nsText.substring(with: NSRange(location: location, length: exactLength)),
            suffix: suffixStart < length
                ? nsText.substring(with: NSRange(location: suffixStart, length: min(8, length - suffixStart)))
                : nil
        )
    }

    private static func makeRowColumnIndices(for text: String) -> [RDEPUBRowColumnIndex] {
        let nsText = text as NSString
        let length = nsText.length
        guard length > 0 else {
            return [RDEPUBRowColumnIndex(row: 0, startOffset: 0, endOffset: 0)]
        }

        var rows: [RDEPUBRowColumnIndex] = []
        var rowNumber = 0
        var lineStart = 0

        nsText.enumerateSubstrings(
            in: NSRange(location: 0, length: length),
            options: [.byLines, .substringNotRequired]
        ) { _, substringRange, enclosingRange, _ in
            let startOffset = enclosingRange.location
            let lineLength = max(substringRange.length, 0)
            let endOffset = max(startOffset + max(lineLength - 1, 0), startOffset)
            rows.append(
                RDEPUBRowColumnIndex(
                    row: rowNumber,
                    startOffset: startOffset,
                    endOffset: min(endOffset, max(length - 1, 0))
                )
            )
            rowNumber += 1
            lineStart = enclosingRange.location + enclosingRange.length
        }

        if rows.isEmpty {
            rows.append(RDEPUBRowColumnIndex(row: 0, startOffset: 0, endOffset: max(length - 1, 0)))
        } else if lineStart == length, text.hasSuffix("\n") {
            rows.append(RDEPUBRowColumnIndex(row: rowNumber, startOffset: length, endOffset: length))
        }

        return rows
    }
}