ReadViewSDK/Sources/RDReaderView/EPUBUI/ReaderController/RDEPUBReaderSearchCoordinator.swift

import Foundation

final class RDEPUBReaderSearchCoordinator {

    private unowned let context: RDEPUBReaderContext

    private let searchQueue = DispatchQueue(label: "com.ssreaderview.epub.search", qos: .userInitiated)

    private let tokenLock = NSLock()
    private var _currentSearchToken: UUID = UUID()

    private var currentSearchToken: UUID {
        get {
            tokenLock.lock()
            defer { tokenLock.unlock() }
            return _currentSearchToken
        }
        set {
            tokenLock.lock()
            _currentSearchToken = newValue
            tokenLock.unlock()
        }
    }

    init(context: RDEPUBReaderContext) {
        self.context = context
    }

    private var controller: RDEPUBReaderController? {
        context.controller
    }

    func search(keyword: String) {
        guard let controller else { return }
        let normalizedKeyword = keyword.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !normalizedKeyword.isEmpty else {
            clearSearch()
            return
        }

        let token = UUID()
        currentSearchToken = token

        let searchEngine = makeSearchEngine()
        let layoutSnapshot = context.makeLayoutSnapshot()

        searchQueue.async { [weak self] in
            guard let self, self.currentSearchToken == token else { return }
            let matches = self.performSearch(using: searchEngine, keyword: normalizedKeyword, layoutSnapshot: layoutSnapshot)

            DispatchQueue.main.async { [weak self] in
                guard let self, self.currentSearchToken == token else { return }
                guard let controller = self.context.controller else { return }
                controller.searchState = RDEPUBSearchState(
                    keyword: normalizedKeyword,
                    matches: matches,
                    currentMatchIndex: matches.isEmpty ? nil : 0
                )
                self.notifySearchStateChanged()
                if matches.isEmpty {
                    controller.refreshVisibleContentPreservingLocation()
                } else {
                    _ = self.navigateToCurrentSearchMatch(animated: false)
                }
            }
        }
    }

    @discardableResult
    func searchNext() -> Bool {
        advanceSearch(by: 1)
    }

    @discardableResult
    func searchPrevious() -> Bool {
        advanceSearch(by: -1)
    }

    @discardableResult
    func selectSearchMatch(at index: Int) -> Bool {
        guard let controller else { return false }
        guard var searchState = controller.searchState,
              searchState.matches.indices.contains(index) else {
            return false
        }
        searchState.currentMatchIndex = index
        controller.searchState = searchState
        notifySearchStateChanged()
        return navigateToCurrentSearchMatch(animated: true)
    }

    func clearSearch() {
        currentSearchToken = UUID()
        guard let controller else { return }
        controller.searchState = nil
        notifySearchStateChanged()
        controller.refreshVisibleContentPreservingLocation()
    }

    func searchPresentation(for page: EPUBPage) -> RDEPUBSearchPresentation? {
        guard let controller else { return nil }
        guard let searchState = controller.searchState,
              let publication = controller.publication else {
            return nil
        }

        let pageHrefs: [String]
        if let fixedSpread = page.fixedSpread {
            pageHrefs = fixedSpread.resources.map { publication.resourceResolver.normalizedHref($0.href) ?? $0.href }
        } else if publication.spine.indices.contains(page.spineIndex) {
            pageHrefs = [
                publication.resourceResolver.normalizedHref(publication.spine[page.spineIndex].href)
                    ?? publication.spine[page.spineIndex].href
            ]
        } else {
            pageHrefs = []
        }

        guard !pageHrefs.isEmpty else {
            return nil
        }

        let currentMatch = searchState.currentMatch
        let normalizedCurrentHref = currentMatch.map { publication.resourceResolver.normalizedHref($0.href) ?? $0.href }
        let resources = pageHrefs.map { href in
            let matchCount = searchState.matches.filter {
                (publication.resourceResolver.normalizedHref($0.href) ?? $0.href) == href
            }.count
            let activeLocalMatchIndex = normalizedCurrentHref == href ? currentMatch?.localMatchIndex : nil
            return RDEPUBSearchPresentationResource(
                href: href,
                matchCount: matchCount,
                activeLocalMatchIndex: activeLocalMatchIndex
            )
        }
        return RDEPUBSearchPresentation(keyword: searchState.keyword, resources: resources)
    }

    // MARK: - Private

    private func makeSearchEngine() -> SearchEngineSnapshot? {
        guard let controller else { return nil }
        if let textBook = controller.textBook, let publication = controller.publication {
            return .textBook(textBook, publication)
        }
        // External text books (e.g. plain .txt) carry a textBook but no
        // publication/parser/bookPageMap; search over the chapter text
        // directly using hrefs from the textBook itself.
        if let textBook = controller.textBook {
            return .externalTextBook(textBook)
        }
        if controller.readerContext.bookPageMap != nil, let publication = controller.publication {
            return .onDemand(publication)
        }
        if let parser = controller.parser, let publication = controller.publication {
            return .html(parser, publication)
        }
        return nil
    }

    private func performSearch(using engine: SearchEngineSnapshot?, keyword: String, layoutSnapshot: RDEPUBLayoutSnapshot?) -> [RDEPUBSearchMatch] {
        guard let engine else { return [] }
        switch engine {
        case .textBook(let textBook, let publication):
            return RDEPUBTextSearchEngine(textBook: textBook, publication: publication).search(keyword: keyword)
        case .externalTextBook(let textBook):
            return RDEPUBTextSearchEngine.searchWithoutPublication(textBook: textBook, keyword: keyword)
        case .onDemand(let publication):
            return resolvedOnDemandSearchMatches(for: keyword, publication: publication, layoutSnapshot: layoutSnapshot)
        case .html(let parser, let publication):
            return RDEPUBHTMLSearchEngine(parser: parser, publication: publication).search(keyword: keyword)
        }
    }

    private enum SearchEngineSnapshot {
        case textBook(RDEPUBTextBook, RDEPUBPublication)
        case externalTextBook(RDEPUBTextBook)
        case onDemand(RDEPUBPublication)
        case html(RDEPUBParser, RDEPUBPublication)
    }

    // H-09: Each chapter iteration is wrapped in autoreleasepool to release
    // the chapter's typesetAttributedString memory between iterations.
    private func resolvedOnDemandSearchMatches(for keyword: String, publication: RDEPUBPublication, layoutSnapshot: RDEPUBLayoutSnapshot?) -> [RDEPUBSearchMatch] {
        let normalizedKeyword = keyword.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !normalizedKeyword.isEmpty else {
            return []
        }

        let buildableSpineIndices = publication.spine.indices.filter { index in
            let item = publication.spine[index]
            return item.linear && (item.mediaType.contains("html") || item.mediaType.contains("xhtml"))
        }

        var matches: [RDEPUBSearchMatch] = []
        for spineIndex in buildableSpineIndices {
            let chapterMatches: [RDEPUBSearchMatch] = autoreleasepool {
                guard let chapter = try? context.runtime?.chapterLoader.loadChapterSynchronouslyForMigration(
                    spineIndex: spineIndex,
                    store: context.runtime?.chapterRuntimeStore ?? RDEPUBChapterRuntimeStore(),
                    layoutSnapshot: layoutSnapshot
                ) else {
                    return []
                }

                let chapterData = makeChapterData(from: chapter, chapterIndex: chapter.pages.first?.chapterIndex ?? 0)
                let source = chapter.typesetAttributedString.string as NSString
                let fullLength = source.length
                guard fullLength > 0 else { return [] }

                let normalizedHref = publication.resourceResolver.normalizedHref(chapter.href) ?? chapter.href
                var localMatches: [RDEPUBSearchMatch] = []
                var localMatchIndex = 0
                var searchRange = NSRange(location: 0, length: fullLength)

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

                    let progressionDenominator = max(fullLength - 1, 1)
                    let progression = Double(foundRange.location) / Double(progressionDenominator)
                    let rangeAnchor = chapterData.rangeAnchor(for: foundRange)
                    localMatches.append(
                        RDEPUBSearchMatch(
                            href: normalizedHref,
                            progression: progression,
                            previewText: previewText(in: source, matchRange: foundRange),
                            localMatchIndex: localMatchIndex,
                            rangeLocation: foundRange.location,
                            rangeLength: max(foundRange.length, 1),
                            rangeAnchor: rangeAnchor,
                            cfi: chapterData.indexTable.cfi(for: rangeAnchor.start)?.rawValue,
                            rangeCFI: chapterData.indexTable.cfiRange(for: rangeAnchor)?.rawValue
                        )
                    )

                    localMatchIndex += 1
                    let nextLocation = foundRange.location + max(foundRange.length, 1)
                    if nextLocation >= fullLength {
                        break
                    }
                    searchRange = NSRange(location: nextLocation, length: fullLength - nextLocation)
                }
                return localMatches
            } // end autoreleasepool
            matches.append(contentsOf: chapterMatches)
        }

        return matches
    }

    private func makeChapterData(
        from runtimeChapter: RDEPUBRuntimeChapter,
        chapterIndex: Int
    ) -> RDEPUBChapterData {
        let textChapter = RDEPUBTextChapter(
            chapterIndex: chapterIndex,
            spineIndex: runtimeChapter.spineIndex,
            href: runtimeChapter.href,
            title: runtimeChapter.title,
            attributedContent: runtimeChapter.typesetAttributedString,
            fragmentOffsets: runtimeChapter.chapterOffsetMap.fragmentOffsets,
            cfiMap: runtimeChapter.chapterOffsetMap.cfiMap,
            pageBreakReasons: runtimeChapter.pages.map(\.metadata.breakReason),
            pages: runtimeChapter.pages
        )
        return RDEPUBChapterData(
            chapter: textChapter,
            indexTable: RDEPUBTextIndexTable(chapters: [textChapter])
        )
    }

    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)
    }

    private func advanceSearch(by delta: Int) -> Bool {
        guard let controller else { return false }
        guard var searchState = controller.searchState, !searchState.matches.isEmpty else {
            return false
        }

        let currentIndex = searchState.currentMatchIndex ?? 0
        let nextIndex = (currentIndex + delta + searchState.matches.count) % searchState.matches.count
        searchState.currentMatchIndex = nextIndex
        controller.searchState = searchState
        notifySearchStateChanged()
        return navigateToCurrentSearchMatch(animated: true)
    }

    private func notifySearchStateChanged() {
        guard let controller else { return }
        let state = controller.searchState
        controller.delegate?.epubReader(controller, didUpdateSearchResult: state?.result)
        controller.delegate?.epubReader(controller, didChangeCurrentSearchMatch: state?.currentMatch)
    }

    private func navigateToCurrentSearchMatch(animated: Bool) -> Bool {
        guard let controller else { return false }
        guard let searchMatch = controller.searchState?.currentMatch else {
            controller.refreshVisibleContentPreservingLocation()
            return false
        }

        if let targetPageNumber = pageNumber(for: searchMatch),
           controller.readerView.currentPage == targetPageNumber - 1 {
            controller.refreshVisibleContentPreservingLocation()
            return true
        }

        let location = RDEPUBLocation(
            bookIdentifier: controller.currentBookIdentifier,
            href: searchMatch.href,
            progression: searchMatch.progression,
            lastProgression: searchMatch.progression,
            fragment: nil,
            rangeAnchor: searchMatch.rangeAnchor,
            cfi: searchMatch.cfi,
            rangeCFI: searchMatch.rangeCFI
        )
        return controller.restoreReadingLocation(location, animated: animated)
    }

    private func pageNumber(for searchMatch: RDEPUBSearchMatch) -> Int? {
        guard let controller else { return nil }
        if let chapterData = controller.textChapterData(forNormalizedHref: searchMatch.href) {
            if let exactPageNumber = exactPageNumber(
                for: searchMatch,
                in: chapterData,
                keyword: controller.searchState?.keyword
            ) {
                return exactPageNumber
            }

            if let pageNumber = chapterData.pageNumber(for: searchMatch) {
                return pageNumber
            }

            if let rangeLocation = searchMatch.rangeLocation,
               let page = chapterData.page(containing: rangeLocation) {
                return page.absolutePageIndex + 1
            }
        }

        let location = RDEPUBLocation(
            bookIdentifier: controller.currentBookIdentifier,
            href: searchMatch.href,
            progression: searchMatch.progression,
            lastProgression: searchMatch.progression,
            fragment: nil,
            rangeAnchor: searchMatch.rangeAnchor,
            cfi: searchMatch.cfi,
            rangeCFI: searchMatch.rangeCFI
        )

        if let textBook = controller.textBook {
            if let publication = controller.publication {
                return textBook.pageNumber(
                    for: location,
                    resolver: publication.resourceResolver,
                    bookIdentifier: controller.currentBookIdentifier
                )
            }
            // External text books have no resolver; hrefs match verbatim.
            return textBook.chapterData(for: location.href)?.pageNumber(for: location)
        }

        return controller.readingSession?.pageIndex(
            for: location,
            bookIdentifier: controller.currentBookIdentifier
        ).map { $0 + 1 }
    }

    private func exactPageNumber(
        for searchMatch: RDEPUBSearchMatch,
        in chapterData: RDEPUBChapterData,
        keyword: String?
    ) -> Int? {
        let normalizedKeyword = keyword?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
        guard !normalizedKeyword.isEmpty else { return nil }

        let source = chapterData.attributedContent.string as NSString
        let fullLength = source.length
        guard fullLength > 0 else { return nil }

        var localMatchIndex = 0
        var searchRange = NSRange(location: 0, length: fullLength)

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

            if localMatchIndex == searchMatch.localMatchIndex,
               let page = chapterData.page(containing: foundRange.location) {
                return page.absolutePageIndex + 1
            }

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

        return nil
    }
}