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

import Foundation

public struct RDEPUBCFIRecoveryResult: Equatable {
    public enum Confidence: Int, Codable {
        case exactPath
        case assertionCalibrated
        case siblingRecovered
        case tokenRecovered
        case fragmentFallback
        case offsetFallback
    }

    public var chapterOffset: Int
    public var confidence: Confidence

    public init(chapterOffset: Int, confidence: Confidence) {
        self.chapterOffset = chapterOffset
        self.confidence = confidence
    }
}

enum RDEPUBCFIRecoveryEngine {
    static func recover(
        cfi: RDEPUBCFI,
        cfiMap: RDEPUBCFIMap?,
        chapterText: String?,
        fragmentOffsets: [String: Int],
        fallbackOffset: Int?,
        lastOffset: Int
    ) -> RDEPUBCFIRecoveryResult? {
        let resolved = RDEPUBCFIResolver.resolve(cfi)

        if let exact = exactPathResult(cfi: cfi, cfiMap: cfiMap, lastOffset: lastOffset) {
            return calibrateIfNeeded(
                exact,
                cfi: cfi,
                chapterText: chapterText,
                lastOffset: lastOffset
            )
        }

        if let sibling = siblingRecoveredResult(cfi: cfi, cfiMap: cfiMap, lastOffset: lastOffset) {
            return calibrateIfNeeded(
                sibling,
                cfi: cfi,
                chapterText: chapterText,
                lastOffset: lastOffset
            )
        }

        if let token = tokenRecoveredResult(cfi: cfi, cfiMap: cfiMap, chapterText: chapterText, lastOffset: lastOffset) {
            return token
        }

        if let fragmentID = resolved.fragmentID,
           let fragmentOffset = fragmentOffsets[fragmentID]
                ?? cfiMap?.markers.first(where: { $0.fragmentID == fragmentID })?.chapterOffset {
            return RDEPUBCFIRecoveryResult(
                chapterOffset: clamp(fragmentOffset, lastOffset: lastOffset),
                confidence: .fragmentFallback
            )
        }

        if let fallbackOffset {
            return RDEPUBCFIRecoveryResult(
                chapterOffset: clamp(fallbackOffset, lastOffset: lastOffset),
                confidence: .offsetFallback
            )
        }

        return nil
    }

    private static func exactPathResult(
        cfi: RDEPUBCFI,
        cfiMap: RDEPUBCFIMap?,
        lastOffset: Int
    ) -> RDEPUBCFIRecoveryResult? {
        guard let cfiMap else { return nil }
        if let marker = cfiMap.marker(matching: cfi.contentPath),
           let chapterOffset = chapterOffset(for: cfi, marker: marker, lastOffset: lastOffset) {
            return RDEPUBCFIRecoveryResult(chapterOffset: chapterOffset, confidence: .exactPath)
        }
        if let pathRange = cfiMap.pathRanges.first(where: { $0.cfiPath == cfi.contentPath }) {
            let localOffset = max(cfi.characterOffset ?? 0, 0)
            let chapterOffset = clamp(pathRange.startOffset + min(localOffset, max(pathRange.textNodeLength - 1, 0)), lastOffset: lastOffset)
            return RDEPUBCFIRecoveryResult(chapterOffset: chapterOffset, confidence: .exactPath)
        }
        return nil
    }

    private static func siblingRecoveredResult(
        cfi: RDEPUBCFI,
        cfiMap: RDEPUBCFIMap?,
        lastOffset: Int
    ) -> RDEPUBCFIRecoveryResult? {
        guard let cfiMap,
              let targetStep = cfi.contentPath.steps.last else { return nil }

        let targetParent = cfi.contentPath.steps.dropLast()
        let candidates = cfiMap.markers.filter { marker in
            marker.cfiPath.steps.count == cfi.contentPath.steps.count
                && Array(marker.cfiPath.steps.dropLast()) == Array(targetParent)
                && marker.chapterOffset != nil
        }
        guard !candidates.isEmpty else { return nil }

        let siblingSignature = siblingSignature(for: cfi.contentPath)
        let best = candidates.min { lhs, rhs in
            siblingScore(for: lhs, targetStep: targetStep, siblingSignature: siblingSignature)
                < siblingScore(for: rhs, targetStep: targetStep, siblingSignature: siblingSignature)
        }

        guard let best,
              let chapterOffset = chapterOffset(for: cfi, marker: best, lastOffset: lastOffset) else {
            return nil
        }

        return RDEPUBCFIRecoveryResult(chapterOffset: chapterOffset, confidence: .siblingRecovered)
    }

    private static func tokenRecoveredResult(
        cfi: RDEPUBCFI,
        cfiMap: RDEPUBCFIMap?,
        chapterText: String?,
        lastOffset: Int
    ) -> RDEPUBCFIRecoveryResult? {
        guard let cfiMap,
              let chapterText,
              !chapterText.isEmpty,
              let exact = cfi.textAssertion?.exact,
              !exact.isEmpty else {
            return nil
        }

        let minLength = 2
        let exactCount = Double(max(exact.count, 1))
        let candidateAnchors = cfiMap.recoveryMetadata.tokenIndex.filter { anchor in
            let tokenCount = anchor.token.count
            guard tokenCount >= minLength else { return false }
            let ratio = Double(tokenCount) / exactCount
            guard ratio >= 0.3 && ratio <= 3.0 else { return false }
            let hasPrefixMatch = anchor.token.hasPrefix(exact) || exact.hasPrefix(anchor.token)
            let hasSuffixMatch = anchor.token.hasSuffix(exact) || exact.hasSuffix(anchor.token)
            return hasPrefixMatch || hasSuffixMatch
        }
        guard !candidateAnchors.isEmpty else { return nil }

        let normalizedText = RDEPUBCFITextNodeMapBuilder.normalizedText(from: chapterText)
        guard !normalizedText.isEmpty else { return nil }

        var bestOffset: Int?
        var bestScore = Int.max

        for anchor in candidateAnchors {
            let approximateOffset = clamp(anchor.chapterOffset, lastOffset: lastOffset)
            let calibrated = calibratedOffset(
                approximateOffset,
                assertion: cfi.textAssertion,
                text: chapterText,
                lastOffset: lastOffset,
                windowRadius: 768
            )
            let score = abs(calibrated - approximateOffset)
            if score < bestScore {
                bestScore = score
                bestOffset = calibrated
            }
        }

        guard let bestOffset else { return nil }
        return RDEPUBCFIRecoveryResult(chapterOffset: bestOffset, confidence: .tokenRecovered)
    }

    private static func calibrateIfNeeded(
        _ result: RDEPUBCFIRecoveryResult,
        cfi: RDEPUBCFI,
        chapterText: String?,
        lastOffset: Int
    ) -> RDEPUBCFIRecoveryResult {
        guard let chapterText,
              let assertion = cfi.textAssertion,
              assertion.exact?.isEmpty == false else {
            return result
        }
        let calibrated = calibratedOffset(
            result.chapterOffset,
            assertion: assertion,
            text: chapterText,
            lastOffset: lastOffset
        )
        if calibrated != result.chapterOffset {
            return RDEPUBCFIRecoveryResult(chapterOffset: calibrated, confidence: .assertionCalibrated)
        }
        return result
    }

    private static func chapterOffset(
        for cfi: RDEPUBCFI,
        marker: RDEPUBCFIMarker,
        lastOffset: Int
    ) -> Int? {
        guard let baseOffset = marker.chapterOffset else { return nil }
        let localOffset = max(cfi.characterOffset ?? 0, 0)
        if let textNodeLength = marker.textNodeLength, textNodeLength > 0 {
            return clamp(baseOffset + min(localOffset, max(textNodeLength - 1, 0)), lastOffset: lastOffset)
        }
        return clamp(baseOffset, lastOffset: lastOffset)
    }

    private static func siblingSignature(for path: RDEPUBCFIPath) -> String {
        guard let last = path.steps.last else { return "" }
        let parent = path.steps.dropLast().map { String($0.index) }.joined(separator: "/")
        let childIndex = max(last.index / 2, 0)
        return "\(parent)#\(childIndex)"
    }

    private static func siblingScore(
        for marker: RDEPUBCFIMarker,
        targetStep: RDEPUBCFIStep,
        siblingSignature: String
    ) -> Int {
        guard let candidateStep = marker.cfiPath.steps.last else { return Int.max }
        var score = abs(candidateStep.index - targetStep.index)
        if marker.domSiblingSignature != siblingSignature {
            score += 1_000
        }
        if marker.fragmentID != nil {
            score += 100
        }
        return score
    }

    private static func calibratedOffset(
        _ offset: Int,
        assertion: RDEPUBCFITextAssertion?,
        text: String,
        lastOffset: Int,
        windowRadius: Int = 512
    ) -> Int {
        let clampedOffset = clamp(offset, lastOffset: lastOffset)
        guard let assertion,
              let exact = assertion.exact,
              !exact.isEmpty else {
            return clampedOffset
        }

        let nsText = text as NSString
        let length = nsText.length
        guard length > 0 else { return clampedOffset }

        let searchStart = max(clampedOffset - windowRadius, 0)
        let searchEnd = min(clampedOffset + windowRadius + exact.utf16.count, length)
        guard searchEnd > searchStart else { return clampedOffset }

        let searchRange = NSRange(location: searchStart, length: searchEnd - searchStart)
        var candidateRange = nsText.range(of: exact, options: [], range: searchRange)
        var bestOffset: Int?
        var bestScore = Int.max

        while candidateRange.location != NSNotFound {
            let score = assertionScore(
                assertion,
                candidateLocation: candidateRange.location,
                exactLength: candidateRange.length,
                in: nsText,
                preferredOffset: clampedOffset
            )
            if score < bestScore {
                bestScore = score
                bestOffset = candidateRange.location
            }

            let nextLocation = candidateRange.location + max(candidateRange.length, 1)
            let upperBound = searchRange.location + searchRange.length
            guard nextLocation < upperBound else { break }
            candidateRange = nsText.range(
                of: exact,
                options: [],
                range: NSRange(location: nextLocation, length: upperBound - nextLocation)
            )
        }

        guard let bestOffset, bestScore < Int.max else {
            return clampedOffset
        }
        return clamp(bestOffset, lastOffset: lastOffset)
    }

    private static func assertionScore(
        _ assertion: RDEPUBCFITextAssertion,
        candidateLocation: Int,
        exactLength: Int,
        in text: NSString,
        preferredOffset: Int
    ) -> Int {
        var score = abs(candidateLocation - preferredOffset)
        if let prefix = assertion.prefix, !prefix.isEmpty {
            let prefixStart = max(candidateLocation - prefix.utf16.count, 0)
            let availableLength = candidateLocation - prefixStart
            let candidatePrefix = availableLength > 0
                ? text.substring(with: NSRange(location: prefixStart, length: availableLength))
                : ""
            score += candidatePrefix.hasSuffix(prefix) ? 0 : 10_000
        }
        if let suffix = assertion.suffix, !suffix.isEmpty {
            let suffixStart = min(candidateLocation + exactLength, text.length)
            let suffixLength = min(suffix.utf16.count, text.length - suffixStart)
            let candidateSuffix = suffixLength > 0
                ? text.substring(with: NSRange(location: suffixStart, length: suffixLength))
                : ""
            score += candidateSuffix == suffix ? 0 : 10_000
        }
        return score
    }

    private static func clamp(_ offset: Int, lastOffset: Int) -> Int {
        min(max(offset, 0), max(lastOffset, 0))
    }
}