ReadViewSDK/Sources/RDReaderView/EPUBCore/CFI/RDEPUBCFIDOMPathBuilder.swift

import Foundation

public enum RDEPUBCFIDOMPathBuilder {
    public static func fragmentPaths(in html: String) -> [String: RDEPUBCFIPath] {
        // Pre-strip HTML comments and CDATA to avoid false tag matches inside them
        let cleaned = stripCommentsAndCDATA(html)

        guard let regex = try? NSRegularExpression(
            pattern: #"</?\s*([A-Za-z][A-Za-z0-9:_-]*)([^>]*)>"#,
            options: [.caseInsensitive]
        ) else {
            return [:]
        }

        let nsHTML = cleaned as NSString
        var stack: [RDEPUBCFIStep] = []
        var childCountsByDepth: [Int: Int] = [:]
        var paths: [String: RDEPUBCFIPath] = [:]

        for match in regex.matches(in: cleaned, range: NSRange(location: 0, length: nsHTML.length)) {
            let rawTag = nsHTML.substring(with: match.range)
            guard match.numberOfRanges > 2 else { continue }
            let tagName = nsHTML.substring(with: match.range(at: 1)).lowercased()
            guard !isIgnorableTag(tagName) else { continue }

            if rawTag.hasPrefix("</") {
                if tagName == "html" {
                    continue
                }
                if !stack.isEmpty {
                    stack.removeLast()
                    childCountsByDepth[stack.count + 1] = nil
                }
                continue
            }

            if tagName == "html" {
                childCountsByDepth[0] = nil
                continue
            }

            let attributes = nsHTML.substring(with: match.range(at: 2))
            let depth = stack.count
            let nextChildIndex = (childCountsByDepth[depth] ?? 0) + 1
            childCountsByDepth[depth] = nextChildIndex

            let fragmentID = idAttribute(in: attributes)
            let step = RDEPUBCFIStep(index: nextChildIndex * 2, idAssertion: fragmentID)
            let currentPath = RDEPUBCFIPath(steps: stack + [step])
            if let fragmentID, paths[fragmentID] == nil {
                paths[fragmentID] = currentPath
            }

            if !rawTag.hasSuffix("/>"), !isVoidTag(tagName) {
                stack.append(step)
                childCountsByDepth[stack.count] = nil
            }
        }

        return paths
    }

    static func idAttribute(in attributes: String) -> String? {
        guard let regex = try? NSRegularExpression(
            pattern: #"(?:^|\s)(?:id|xml:id)\s*=\s*(['"])(.*?)\1"#,
            options: [.caseInsensitive]
        ) else {
            return nil
        }
        let nsAttributes = attributes as NSString
        guard let match = regex.firstMatch(
            in: attributes,
            range: NSRange(location: 0, length: nsAttributes.length)
        ),
              match.numberOfRanges > 2 else {
            return nil
        }
        return nsAttributes.substring(with: match.range(at: 2)).nilIfEmpty
    }

    static func isIgnorableTag(_ tagName: String) -> Bool {
        tagName == "!doctype" || tagName.hasPrefix("?")
    }

    /// Strip HTML comments `<!--...-->` and CDATA sections `<![CDATA[...]]>` from HTML
    /// to prevent the tag regex from matching false tags inside them.
    static func stripCommentsAndCDATA(_ html: String) -> String {
        var result = html
        // Strip comments: <!--...-->
        if let commentRegex = try? NSRegularExpression(pattern: #"<!--[\s\S]*?-->"#, options: []) {
            let nsResult = result as NSString
            let matches = commentRegex.matches(in: result, range: NSRange(location: 0, length: nsResult.length))
            for match in matches.reversed() {
                result = nsResult.replacingCharacters(in: match.range, with: "")
            }
        }
        // Strip CDATA: <![CDATA[...]]>
        if let cdataRegex = try? NSRegularExpression(pattern: #"<!\[CDATA\[[\s\S]*?\]\]>"#, options: []) {
            let nsResult = result as NSString
            let matches = cdataRegex.matches(in: result, range: NSRange(location: 0, length: nsResult.length))
            for match in matches.reversed() {
                result = nsResult.replacingCharacters(in: match.range, with: "")
            }
        }
        return result
    }

    static func isVoidTag(_ tagName: String) -> Bool {
        [
            "area", "base", "br", "col", "embed", "hr", "img", "input",
            "link", "meta", "param", "source", "track", "wbr"
        ].contains(tagName)
    }
}

enum RDEPUBCFITextNodeMapBuilder {
    static func normalizedText(from source: String) -> String {
        RDEPUBNormalizedTextIndex.normalize(source)
    }

    static func makeMap(
        href: String,
        rawHTML: String,
        chapterText: String,
        fragmentOffsets: [String: Int]
    ) -> RDEPUBCFIMap {
        let normalizedChapter = RDEPUBNormalizedTextIndex(source: chapterText)
        let fragmentPaths = RDEPUBCFIDOMPathBuilder.fragmentPaths(in: rawHTML)
        let markers = makeMarkers(
            href: href,
            rawHTML: rawHTML,
            chapterText: chapterText,
            fragmentOffsets: fragmentOffsets
        )
        let pathRanges = makePathRanges(from: markers)
        return RDEPUBCFIMap(
            href: href,
            markers: markers,
            pathRanges: pathRanges,
            recoveryMetadata: RDEPUBCFIRecoveryMetadata(
                domFingerprint: rawHTML.sha256Hex,
                normalizedTextChecksum: normalizedChapter.normalizedText.sha256Hex,
                tokenIndex: makeTokenAnchors(
                    normalizedChapter: normalizedChapter,
                    markers: markers,
                    pathRanges: pathRanges
                ),
                fragmentPathMap: fragmentPaths
            )
        )
    }

    static func makeMarkers(
        href: String,
        rawHTML: String,
        chapterText: String,
        fragmentOffsets: [String: Int]
    ) -> [RDEPUBCFIMarker] {
        let normalizedChapter = RDEPUBNormalizedTextIndex(source: chapterText)
        guard !normalizedChapter.normalizedText.isEmpty else {
            return fragmentMarkers(fragmentOffsets: fragmentOffsets)
        }

        // Use comment/CDATA-stripped HTML for both path building and tag matching
        let cleanedHTML = stripCommentsAndCDATA(rawHTML)

        let fragments = RDEPUBCFIDOMPathBuilder.fragmentPaths(in: rawHTML)
        var markers: [RDEPUBCFIMarker] = fragmentMarkers(fragmentOffsets: fragmentOffsets, paths: fragments)
        var searchLocation = 0
        var stack: [RDEPUBCFIStep] = []
        var childCountsByDepth: [Int: Int] = [:]
        var lastTagEnd = cleanedHTML.startIndex

        for tagMatch in RDEPUBHTMLTagMatch.matches(in: cleanedHTML) {
            let textChunk = String(cleanedHTML[lastTagEnd..<tagMatch.range.lowerBound])
            if let marker = textMarker(
                for: textChunk,
                stack: stack,
                childCountsByDepth: &childCountsByDepth,
                normalizedChapter: normalizedChapter,
                searchLocation: &searchLocation
            ) {
                markers.append(marker)
            }

            processTag(
                tagMatch,
                stack: &stack,
                childCountsByDepth: &childCountsByDepth
            )
            lastTagEnd = tagMatch.range.upperBound
        }

        if lastTagEnd < cleanedHTML.endIndex {
            let trailingText = String(cleanedHTML[lastTagEnd..<cleanedHTML.endIndex])
            if let marker = textMarker(
                for: trailingText,
                stack: stack,
                childCountsByDepth: &childCountsByDepth,
                normalizedChapter: normalizedChapter,
                searchLocation: &searchLocation
            ) {
                markers.append(marker)
            }
        }

        return markers
            .filter { $0.chapterOffset != nil }
            .sorted { lhs, rhs in
                let leftOffset = lhs.chapterOffset ?? Int.max
                let rightOffset = rhs.chapterOffset ?? Int.max
                if leftOffset != rightOffset {
                    return leftOffset < rightOffset
                }
                return lhs.cfiPath.steps.count < rhs.cfiPath.steps.count
            }
    }

    private static func fragmentMarkers(
        fragmentOffsets: [String: Int],
        paths: [String: RDEPUBCFIPath] = [:]
    ) -> [RDEPUBCFIMarker] {
        fragmentOffsets
            .sorted { $0.value < $1.value }
            .map { fragmentID, offset in
                RDEPUBCFIMarker(
                    cfiPath: paths[fragmentID] ?? RDEPUBCFIPath(),
                    chapterOffset: offset,
                    fragmentID: fragmentID
                )
            }
    }

    private static func textMarker(
        for rawText: String,
        stack: [RDEPUBCFIStep],
        childCountsByDepth: inout [Int: Int],
        normalizedChapter: RDEPUBNormalizedTextIndex,
        searchLocation: inout Int
    ) -> RDEPUBCFIMarker? {
        let normalizedNodeText = RDEPUBNormalizedTextIndex.normalize(rawText)
        guard !normalizedNodeText.isEmpty else { return nil }

        let depth = stack.count
        let nextChildIndex = (childCountsByDepth[depth] ?? 0) + 1
        childCountsByDepth[depth] = nextChildIndex
        let textStep = RDEPUBCFIStep(index: nextChildIndex * 2 - 1)
        let path = RDEPUBCFIPath(steps: stack + [textStep])

        // Search using UTF-16 offsets to match normalizedToChapterOffsets indices
        let nsNormalized = normalizedChapter.normalizedText as NSString
        let searchRange = NSRange(location: searchLocation, length: nsNormalized.length - searchLocation)
        guard let foundRange = nsNormalized.range(of: normalizedNodeText, options: [], range: searchRange).toRange(),
              foundRange.lowerBound >= searchLocation else {
            return nil
        }

        searchLocation = foundRange.upperBound
        let normalizedOffset = foundRange.lowerBound
        guard let chapterOffset = normalizedChapter.chapterOffset(forNormalizedOffset: normalizedOffset) else {
            return nil
        }

        return RDEPUBCFIMarker(
            cfiPath: path,
            chapterOffset: chapterOffset,
            textNodeLength: normalizedNodeText.utf16.count,
            textNodeChecksum: normalizedNodeText.fnv1a64,
            normalizedTextPreview: previewText(for: normalizedNodeText),
            domSiblingSignature: siblingSignature(stack: stack, childIndex: nextChildIndex)
        )
    }

    private static func processTag(
        _ match: RDEPUBHTMLTagMatch,
        stack: inout [RDEPUBCFIStep],
        childCountsByDepth: inout [Int: Int]
    ) {
        let tagName = match.name.lowercased()
        guard !RDEPUBCFIDOMPathBuilder.isIgnorableTag(tagName) else { return }

        if match.isClosing {
            if tagName == "html" {
                return
            }
            if !stack.isEmpty {
                stack.removeLast()
                childCountsByDepth[stack.count + 1] = nil
            }
            return
        }

        if tagName == "html" {
            childCountsByDepth[0] = nil
            return
        }

        let depth = stack.count
        let nextChildIndex = (childCountsByDepth[depth] ?? 0) + 1
        childCountsByDepth[depth] = nextChildIndex

        let fragmentID = RDEPUBCFIDOMPathBuilder.idAttribute(in: match.attributes)
        let step = RDEPUBCFIStep(index: nextChildIndex * 2, idAssertion: fragmentID)
        if !match.isSelfClosing && !RDEPUBCFIDOMPathBuilder.isVoidTag(tagName) {
            stack.append(step)
            childCountsByDepth[stack.count] = nil
        }
    }

    private static func makePathRanges(from markers: [RDEPUBCFIMarker]) -> [RDEPUBCFIPathRange] {
        markers.compactMap { marker in
            guard let startOffset = marker.chapterOffset,
                  let textNodeLength = marker.textNodeLength,
                  textNodeLength > 0 else {
                return nil
            }
            return RDEPUBCFIPathRange(
                cfiPath: marker.cfiPath,
                startOffset: startOffset,
                endOffset: startOffset + max(textNodeLength - 1, 0),
                textNodeLength: textNodeLength
            )
        }
    }

    private static func makeTokenAnchors(
        normalizedChapter: RDEPUBNormalizedTextIndex,
        markers: [RDEPUBCFIMarker],
        pathRanges: [RDEPUBCFIPathRange]
    ) -> [RDEPUBCFITokenAnchor] {
        guard !normalizedChapter.normalizedText.isEmpty else { return [] }

        let sortedRanges = pathRanges.sorted { $0.startOffset < $1.startOffset }
        var occurrenceByToken: [String: Int] = [:]
        var anchors: [RDEPUBCFITokenAnchor] = []
        let samples = normalizedChapter.tokenSamples(maxSamples: 64, stride: 96, window: 12)

        for sample in samples {
            guard let chapterOffset = normalizedChapter.chapterOffset(forNormalizedOffset: sample.offset),
                  let path = path(for: chapterOffset, ranges: sortedRanges, markers: markers) else {
                continue
            }
            let occurrence = (occurrenceByToken[sample.token] ?? 0) + 1
            occurrenceByToken[sample.token] = occurrence
            anchors.append(
                RDEPUBCFITokenAnchor(
                    token: sample.token,
                    occurrence: occurrence,
                    chapterOffset: chapterOffset,
                    cfiPath: path
                )
            )
        }

        return anchors
    }

    private static func path(
        for chapterOffset: Int,
        ranges: [RDEPUBCFIPathRange],
        markers: [RDEPUBCFIMarker]
    ) -> RDEPUBCFIPath? {
        if let range = ranges.first(where: { chapterOffset >= $0.startOffset && chapterOffset <= $0.endOffset }) {
            return range.cfiPath
        }
        return markers
            .filter { $0.chapterOffset != nil }
            .min(by: {
                abs(($0.chapterOffset ?? 0) - chapterOffset) < abs(($1.chapterOffset ?? 0) - chapterOffset)
            })?
            .cfiPath
    }

    private static func previewText(for normalizedText: String, limit: Int = 24) -> String {
        String(normalizedText.prefix(limit))
    }

    private static func siblingSignature(stack: [RDEPUBCFIStep], childIndex: Int) -> String {
        let parent = stack.map { String($0.index) }.joined(separator: "/")
        return "\(parent)#\(childIndex)"
    }
}

private struct RDEPUBHTMLTagMatch {
    let range: Range<String.Index>
    let name: String
    let attributes: String
    let isClosing: Bool
    let isSelfClosing: Bool

    static func matches(in html: String) -> [RDEPUBHTMLTagMatch] {
        let cleaned = RDEPUBCFIDOMPathBuilder.stripCommentsAndCDATA(html)

        guard let regex = try? NSRegularExpression(
            pattern: #"</?\s*([A-Za-z][A-Za-z0-9:_-]*)([^>]*)>"#,
            options: [.caseInsensitive]
        ) else {
            return []
        }

        let nsHTML = cleaned as NSString
        return regex.matches(in: cleaned, range: NSRange(location: 0, length: nsHTML.length)).compactMap { match in
            guard match.numberOfRanges > 2,
                  let range = Range(match.range, in: cleaned),
                  let nameRange = Range(match.range(at: 1), in: cleaned),
                  let attributesRange = Range(match.range(at: 2), in: cleaned) else {
                return nil
            }

            let rawTag = String(cleaned[range])
            return RDEPUBHTMLTagMatch(
                range: range,
                name: String(cleaned[nameRange]),
                attributes: String(cleaned[attributesRange]),
                isClosing: rawTag.hasPrefix("</"),
                isSelfClosing: rawTag.hasSuffix("/>")
            )
        }
    }
}

private struct RDEPUBNormalizedTextIndex {
    let normalizedText: String
    private let normalizedToChapterOffsets: [Int]

    init(source: String) {
        var text = ""
        var offsets: [Int] = []
        let nsSource = source as NSString
        var lastWasWhitespace = false

        for index in 0..<nsSource.length {
            let scalar = nsSource.substring(with: NSRange(location: index, length: 1))
            let normalizedScalar = Self.normalizeCharacter(scalar)
            if normalizedScalar == " " {
                if lastWasWhitespace {
                    continue
                }
                lastWasWhitespace = true
            } else {
                lastWasWhitespace = false
            }
            text.append(normalizedScalar)
            offsets.append(index)
        }

        self.normalizedText = text
        self.normalizedToChapterOffsets = offsets
    }

    func chapterOffset(forNormalizedOffset offset: Int) -> Int? {
        guard normalizedToChapterOffsets.indices.contains(offset) else { return nil }
        return normalizedToChapterOffsets[offset]
    }

    func tokenSamples(maxSamples: Int, stride: Int, window: Int) -> [(token: String, offset: Int)] {
        guard !normalizedText.isEmpty, stride > 0, window > 0 else { return [] }
        // Use NSString-based (UTF-16) iteration to match normalizedToChapterOffsets indices
        let nsText = normalizedText as NSString
        let length = nsText.length
        var samples: [(token: String, offset: Int)] = []
        var index = 0

        while index < length, samples.count < maxSamples {
            let end = min(index + window, length)
            let token = nsText.substring(with: NSRange(location: index, length: end - index))
                .trimmingCharacters(in: .whitespacesAndNewlines)
            if token.utf16.count >= min(window, 4) {
                samples.append((token, index))
            }
            index += stride
        }

        if samples.isEmpty {
            let token = nsText.substring(with: NSRange(location: 0, length: min(window, length)))
                .trimmingCharacters(in: .whitespacesAndNewlines)
            if !token.isEmpty {
                samples.append((token, 0))
            }
        }

        return samples
    }

    static func normalize(_ source: String) -> String {
        var result = ""
        var lastWasWhitespace = false
        let decodedSource = RDEPUBHTMLEntityDecoder.decodeEntities(in: source)
        let nsSource = decodedSource as NSString

        for index in 0..<nsSource.length {
            let scalar = nsSource.substring(with: NSRange(location: index, length: 1))
            let normalizedScalar = normalizeCharacter(scalar)
            if normalizedScalar == " " {
                if lastWasWhitespace {
                    continue
                }
                lastWasWhitespace = true
            } else {
                lastWasWhitespace = false
            }
            result.append(normalizedScalar)
        }

        return result
    }

    private static func normalizeCharacter(_ scalar: String) -> Character {
        if scalar == "\u{00A0}" {
            return " "
        }
        if scalar.unicodeScalars.allSatisfy(\.properties.isWhitespace) {
            return " "
        }
        return scalar.first ?? " "
    }
}

private extension String {
    var fnv1a64: UInt64 {
        let offsetBasis: UInt64 = 14695981039346656037
        let prime: UInt64 = 1099511628211
        return utf8.reduce(offsetBasis) { partialResult, byte in
            (partialResult ^ UInt64(byte)) &* prime
        }
    }
}

private enum RDEPUBHTMLEntityDecoder {
    private static let namedEntities: [String: String] = [
        "&nbsp;": " ",
        "&amp;": "&",
        "&lt;": "<",
        "&gt;": ">",
        "&quot;": "\"",
        "&apos;": "'"
    ]

    static func decode(_ value: String) -> String {
        if let mapped = namedEntities[value] {
            return mapped
        }
        guard value.hasPrefix("&"), value.hasSuffix(";") else {
            return value
        }
        let body = String(value.dropFirst().dropLast())
        if body.hasPrefix("#x") || body.hasPrefix("#X"),
           let scalar = UInt32(body.dropFirst(2), radix: 16).flatMap(UnicodeScalar.init) {
            return String(Character(scalar))
        }
        if body.hasPrefix("#"),
           let scalar = UInt32(body.dropFirst(), radix: 10).flatMap(UnicodeScalar.init) {
            return String(Character(scalar))
        }
        return value
    }

    static func decodeEntities(in value: String) -> String {
        guard value.contains("&"),
              let regex = try? NSRegularExpression(pattern: #"&(?:[A-Za-z]+|#\d+|#x[0-9A-Fa-f]+);"#) else {
            return value
        }

        let nsValue = value as NSString
        var decoded = value
        for match in regex.matches(in: value, range: NSRange(location: 0, length: nsValue.length)).reversed() {
            let entity = nsValue.substring(with: match.range)
            let replacement = decode(entity)
            if let range = Range(match.range, in: decoded) {
                decoded.replaceSubrange(range, with: replacement)
            }
        }
        return decoded
    }
}