ReadViewSDK/Sources/RDEpubReaderView/AI/RDEPUBAIContentProvider.swift

import Foundation
import RDAIReaderView
import UIKit

/// EPUB adapter using normalized hrefs and the existing CFI index table for
/// stable citation navigation across font and pagination changes.
@MainActor
public final class RDEPUBAIContentProvider: RDAIContentProvider {
    private weak var controller: RDEPUBReaderController?

    init(controller: RDEPUBReaderController) {
        self.controller = controller
    }

    public func aiDocumentDescriptor() -> RDAIDocumentDescriptor {
        let identifier = controller?.currentBookIdentifier ?? "rd-epub-reader"
        let chapters = readableChapters()
        let revision = chapters.map { "\($0.href)|\($0.text.utf16.count)" }.joined(separator: "|")
        return RDAIDocumentDescriptor(
            identifier: RDAIDocumentIdentifier(rawValue: identifier),
            title: controller?.title ?? "",
            format: .epub,
            contentRevision: RDAIContentHasher.hash(revision)
        )
    }

    public func aiResources() async throws -> [RDAIResourceDescriptor] {
        readableChapters().enumerated().map { index, chapter in
            RDAIResourceDescriptor(
                identifier: RDAIResourceIdentifier(rawValue: chapter.href),
                title: chapter.title,
                order: index,
                estimatedUTF16Length: chapter.text.utf16.count
            )
        }
    }

    public func aiResourceSnapshot(for identifier: RDAIResourceIdentifier) async throws -> RDAIResourceSnapshot {
        guard let chapter = readableChapters().first(where: { $0.href == identifier.rawValue }) else {
            throw RDAIError.resourceUnavailable(identifier)
        }
        let resource = RDAIResourceDescriptor(
            identifier: identifier,
            title: chapter.title,
            order: chapter.order,
            estimatedUTF16Length: chapter.text.utf16.count
        )
        let anchor = RDAIEPUBAnchor(href: chapter.href, progression: 0)
        return RDAIResourceSnapshot(
            descriptor: resource,
            sourceText: chapter.text,
            sourceHash: RDAIContentHasher.hash(chapter.text),
            locatorRuns: [RDAILocatorRun(
                textRange: RDAITextRange(location: 0, length: chapter.text.utf16.count),
                anchor: .epub(anchor)
            )]
        )
    }

    public func aiNavigate(to locator: RDAILocator, animated: Bool) async throws {
        guard let controller,
              case .epub(let anchor) = locator.anchor else {
            throw RDAIError.staleCitation
        }
        let resolved = resolveLocation(locator: locator, anchor: anchor, controller: controller)
        controller.go(to: resolved)
    }

    public func aiShowCitationHighlight(_ citation: RDAICitation) async throws {
        guard let controller,
              case .epub(let anchor) = citation.locator.anchor else {
            throw RDAIError.staleCitation
        }
        let location = resolveLocation(locator: citation.locator, anchor: anchor, controller: controller)
        let range = citation.locator.textRange
        let decoration = RDEPUBHighlight(
            id: "rdai-citation-\(citation.id)",
            bookIdentifier: controller.currentBookIdentifier,
            location: location,
            text: citation.quote,
            rangeInfo: RDEPUBTextOffsetRangeInfo(href: anchor.href, start: range.location, end: range.upperBound).jsonString(),
            style: .highlight,
            color: "#86D7FF"
        )
        controller.setTransientHighlights([decoration])
        controller.go(to: location)
    }

    public func aiClearCitationHighlight() { controller?.clearTransientHighlights() }

    private func resolveLocation(
        locator: RDAILocator,
        anchor: RDAIEPUBAnchor,
        controller: RDEPUBReaderController
    ) -> RDEPUBLocation {
        let sourceLength = readableChapters().first { $0.href == anchor.href }?.text.utf16.count ?? 1
        let progression = Double(locator.textRange.location) / Double(max(sourceLength - 1, 1))
        guard let chapterData = controller.textChapterData(forNormalizedHref: anchor.href) else {
            return RDEPUBLocation(
                bookIdentifier: controller.currentBookIdentifier,
                href: anchor.href,
                progression: anchor.progression ?? progression,
                cfi: anchor.cfi,
                rangeCFI: anchor.rangeCFI
            )
        }
        let range = NSRange(location: locator.textRange.location, length: max(locator.textRange.length, 1))
        let rangeAnchor = chapterData.rangeAnchor(for: range)
        let cfiRange = chapterData.indexTable.cfiRange(for: rangeAnchor)
        return RDEPUBLocation(
            bookIdentifier: controller.currentBookIdentifier,
            href: anchor.href,
            progression: progression,
            rangeAnchor: rangeAnchor,
            cfi: cfiRange?.start.rawValue,
            rangeCFI: cfiRange?.rawValue
        )
    }

    private func readableChapters() -> [(href: String, title: String?, text: String, order: Int)] {
        guard let controller else { return [] }
        if let textBook = controller.textBook {
            return textBook.chapters.enumerated().map {
                (href: $0.element.href, title: $0.element.title, text: $0.element.attributedContent.string, order: $0.offset)
            }
        }
        guard let publication = controller.publication,
              let parser = controller.parser else {
            return []
        }
        return publication.spine.enumerated().compactMap { index, item in
            guard item.linear,
                  item.mediaType.contains("html") || item.mediaType.contains("xhtml"),
                  let html = parser.htmlString(forRelativePath: item.href) else {
                return nil
            }
            let href = publication.resourceResolver.normalizedHref(item.href) ?? item.href
            return (href: href, title: item.href, text: plainText(fromHTML: html), order: index)
        }
    }

    private func plainText(fromHTML html: String) -> String {
        guard let data = html.data(using: .utf8),
              let attributed = try? NSAttributedString(
                data: data,
                options: [
                    .documentType: NSAttributedString.DocumentType.html,
                    .characterEncoding: String.Encoding.utf8.rawValue
                ],
                documentAttributes: nil
              ) else {
            return html.replacingOccurrences(of: "<[^>]+>", with: " ", options: .regularExpression)
        }
        return attributed.string
    }
}

public extension RDEPUBReaderController {
    func aiCurrentReadScope() -> RDAIReadScope {
        guard let location = currentLocation else { return .init() }
        let locator = RDAILocator(
            documentIdentifier: RDAIDocumentIdentifier(rawValue: currentBookIdentifier ?? "rd-epub-reader"),
            resourceIdentifier: RDAIResourceIdentifier(rawValue: location.href),
            textRange: RDAITextRange(location: 0, length: Int.max),
            anchor: .epub(.init(href: location.href, cfi: location.cfi, rangeCFI: location.rangeCFI, progression: location.progression)),
            sourceHash: ""
        )
        return RDAIReadScope(upperBound: locator)
    }

    func makeAIContentProvider() -> RDEPUBAIContentProvider {
        RDEPUBAIContentProvider(controller: self)
    }

    func makeAIReaderService(
        generativeProvider: (any RDAIGenerativeProvider)? = nil,
        configuration: RDAIReaderConfiguration = .default
    ) -> RDAIReaderService {
        RDAIReaderService(
            contentProvider: makeAIContentProvider(),
            analyzer: RDAINaturalLanguageAnalyzer(),
            semanticScorer: RDAINaturalLanguageSemanticScorer(),
            persistentStore: try? RDAISQLiteIndexStore(),
            generativeProvider: generativeProvider,
            configuration: configuration
        )
    }

    func makeAIReaderAssistant(scope: RDAIReadScope) -> UIViewController {
        UINavigationController(rootViewController: RDAIReaderAssistantViewController(service: makeAIReaderService(), scope: scope))
    }

    func makeAIReaderAssistant(scope: RDAIReadScope, generativeProvider: (any RDAIGenerativeProvider)?) -> UIViewController {
        UINavigationController(rootViewController: RDAIReaderAssistantViewController(service: makeAIReaderService(generativeProvider: generativeProvider), scope: scope))
    }

    func makeAIReaderAssistant() -> UIViewController { makeAIReaderAssistant(scope: aiCurrentReadScope()) }
}