ReadViewSDK/Sources/RDReaderView/EPUBCore/RDEPUBSearchEngine.swift

import Foundation
import UIKit

protocol RDEPUBSearchEngine {

    func search(keyword: String) -> [RDEPUBSearchMatch]
}

final class RDEPUBHTMLSearchEngine: RDEPUBSearchEngine {

    private let parser: RDEPUBParser

    private let publication: RDEPUBPublication

    init(parser: RDEPUBParser, publication: RDEPUBPublication) {
        self.parser = parser
        self.publication = publication
    }

    func search(keyword: String) -> [RDEPUBSearchMatch] {
        let normalizedKeyword = keyword.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !normalizedKeyword.isEmpty else {
            return []
        }

        var searchMatches: [RDEPUBSearchMatch] = []
        for item in publication.spine where item.linear {
            guard item.mediaType.contains("html") || item.mediaType.contains("xhtml"),
                  let html = parser.htmlString(forRelativePath: item.href) else {
                continue
            }

            let plainText = plainText(fromHTML: html, baseURL: parser.fileURL(forRelativePath: item.href)?.deletingLastPathComponent())
            let normalizedHref = publication.resourceResolver.normalizedHref(item.href) ?? item.href
            searchMatches.append(contentsOf: matches(in: plainText, href: normalizedHref, keyword: normalizedKeyword))
        }

        return searchMatches
    }

    private func plainText(fromHTML html: String, baseURL: URL?) -> String {
        guard let data = html.data(using: .utf8) else {
            return fallbackPlainText(fromHTML: html)
        }

        var options: [NSAttributedString.DocumentReadingOptionKey: Any] = [
            .documentType: NSAttributedString.DocumentType.html,
            .characterEncoding: String.Encoding.utf8.rawValue
        ]
        if let baseURL {
            options[NSAttributedString.DocumentReadingOptionKey(rawValue: "NSBaseURLDocumentOption")] = baseURL
        }

        if let attributed = try? NSAttributedString(data: data, options: options, documentAttributes: nil) {
            return attributed.string
        }
        return fallbackPlainText(fromHTML: html)
    }

    private func fallbackPlainText(fromHTML html: String) -> String {
        let stripped = html.replacingOccurrences(of: "<[^>]+>", with: " ", options: .regularExpression)
        return stripped
            .replacingOccurrences(of: "&nbsp;", with: " ")
            .replacingOccurrences(of: "&amp;", with: "&")
            .replacingOccurrences(of: "&lt;", with: "<")
            .replacingOccurrences(of: "&gt;", with: ">")
            .replacingOccurrences(of: "&#39;", with: "'")
            .replacingOccurrences(of: "&quot;", with: "\"")
    }

    private func matches(in text: String, href: String, keyword: String) -> [RDEPUBSearchMatch] {
        let source = text as NSString
        let fullLength = source.length
        guard fullLength > 0 else {
            return []
        }

        var results: [RDEPUBSearchMatch] = []
        var localMatchIndex = 0
        var searchRange = NSRange(location: 0, length: fullLength)

        while searchRange.length > 0 {
            let foundRange = source.range(of: keyword, options: [.caseInsensitive], range: searchRange)
            guard foundRange.location != NSNotFound else {
                break
            }

            let progressionDenominator = max(fullLength - 1, 1)
            let progression = Double(foundRange.location) / Double(progressionDenominator)
            results.append(
                RDEPUBSearchMatch(
                    href: href,
                    progression: progression,
                    previewText: previewText(in: source, matchRange: foundRange),
                    localMatchIndex: localMatchIndex,
                    rangeLocation: foundRange.location,
                    rangeLength: foundRange.length
                )
            )

            localMatchIndex += 1
            let nextLocation = foundRange.location + max(foundRange.length, 1)
            if nextLocation >= fullLength {
                break
            }
            searchRange = NSRange(location: nextLocation, length: fullLength - nextLocation)
        }

        return results
    }

    private func previewText(in text: NSString, matchRange: NSRange) -> String {
        let previewRadius = 12
        let start = max(matchRange.location - previewRadius, 0)
        let end = min(matchRange.location + matchRange.length + previewRadius, text.length)
        let range = NSRange(location: start, length: max(end - start, 0))
        return text.substring(with: range).trimmingCharacters(in: .whitespacesAndNewlines)
    }
}