ReadViewSDK/Sources/RDPDFReaderView/AI/RDPDFAIContentProvider.swift

import CoreGraphics
import Foundation
import RDAIReaderView
import UIKit

/// Bridges PDF text runs into RDAIReaderView without exposing PDF reader UI
/// types to the AI core. The reader remains the source of truth for native
/// text, OCR and page-level highlighting.
@MainActor
public final class RDPDFAIContentProvider: RDAIContentProvider {
    private weak var reader: RDPDFReaderViewController?
    private let book: RDPDFReaderBookDescriptor
    private let extractionCoordinator = RDPDFAIExtractionCoordinator()

    init(reader: RDPDFReaderViewController) {
        self.reader = reader
        book = reader.pageProvider.readerBookDescriptor()
    }

    public func aiDocumentDescriptor() -> RDAIDocumentDescriptor {
        let identifier = RDAIDocumentIdentifier(rawValue: book.identifier)
        return RDAIDocumentDescriptor(
            identifier: identifier,
            title: book.title,
            format: .pdf,
            contentRevision: RDAIContentHasher.hash("\(book.identifier)|\(book.totalPages)")
        )
    }

    public func aiResources() async throws -> [RDAIResourceDescriptor] {
        (0..<book.totalPages).map {
            RDAIResourceDescriptor(identifier: RDAIResourceIdentifier(rawValue: String($0)), order: $0)
        }
    }

    public func aiResourceSnapshot(for identifier: RDAIResourceIdentifier) async throws -> RDAIResourceSnapshot {
        guard let reader,
              let pageIndex = Int(identifier.rawValue),
              pageIndex >= 0,
              pageIndex < book.totalPages else {
            throw RDAIError.resourceUnavailable(identifier)
        }
        let extraction = await extractionCoordinator.text(for: pageIndex, reader: reader)
        let runs = extraction.runs.sorted { $0.readingOrder < $1.readingOrder }
        let textSource = extraction.source
        var sourceText = ""
        var locatorRuns: [RDAILocatorRun] = []
        for run in runs where !run.text.isEmpty {
            if !sourceText.isEmpty { sourceText.append("\n") }
            let range = RDAITextRange(location: sourceText.utf16.count, length: run.text.utf16.count)
            sourceText.append(run.text)
            let source: RDAIPDFAnchor.TextSource = textSource == .ocr ? .ocr : .native
            let anchor = RDAIPDFAnchor(
                pageIndex: pageIndex,
                rects: run.normalizedRects.map(RDAINormalizedRect.init),
                textSource: source,
                readingOrder: run.readingOrder
            )
            locatorRuns.append(RDAILocatorRun(textRange: range, anchor: .pdf(anchor)))
        }
        let descriptor = RDAIResourceDescriptor(
            identifier: identifier,
            order: pageIndex,
            estimatedUTF16Length: sourceText.utf16.count
        )
        return RDAIResourceSnapshot(
            descriptor: descriptor,
            sourceText: sourceText,
            sourceHash: RDAIContentHasher.hash(sourceText),
            locatorRuns: locatorRuns
        )
    }

    public func aiNavigate(to locator: RDAILocator, animated: Bool) async throws {
        guard let reader,
              case .pdf(let anchor) = locator.anchor else {
            throw RDAIError.staleCitation
        }
        reader.showSpeechHighlight(
            pageIndex: anchor.pageIndex,
            normalizedRects: anchor.rects.map(\.cgRect),
            animated: animated
        )
    }

    public func aiShowCitationHighlight(_ citation: RDAICitation) async throws {
        try await aiNavigate(to: citation.locator, animated: true)
    }

    public func aiClearCitationHighlight() {
        reader?.clearSpeechHighlight()
    }
}

private actor RDPDFAIExtractionCoordinator {
    struct Result: Sendable {
        let runs: [RDPDFReaderTextRun]
        let source: RDPDFReaderAnnotationSource
    }

    private var tasks: [Int: Task<Result, Never>] = [:]

    func text(for pageIndex: Int, reader: RDPDFReaderViewController) async -> Result {
        if let task = tasks[pageIndex] { return await task.value }
        let task = Task { @MainActor [weak reader] in
            guard let reader else { return Result(runs: [], source: .region) }
            let runs = await reader.speechTextRuns(at: pageIndex)
            let source = await reader.speechTextSource(at: pageIndex)
            return Result(runs: runs, source: source)
        }
        tasks[pageIndex] = task
        let result = await task.value
        tasks.removeValue(forKey: pageIndex)
        return result
    }
}

public extension RDPDFReaderViewController {
    func aiCurrentReadScope() -> RDAIReadScope {
        let page = currentPageIndex
        let descriptor = pageProvider.readerBookDescriptor()
        let locator = RDAILocator(
            documentIdentifier: RDAIDocumentIdentifier(rawValue: descriptor.identifier),
            resourceIdentifier: RDAIResourceIdentifier(rawValue: String(page)),
            textRange: RDAITextRange(location: 0, length: Int.max),
            anchor: .pdf(.init(pageIndex: page, rects: [], textSource: .native)),
            sourceHash: ""
        )
        return RDAIReadScope(upperBound: locator)
    }

    func makeAIContentProvider() -> RDPDFAIContentProvider {
        RDPDFAIContentProvider(reader: 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()) }
}

private extension RDAINormalizedRect {
    init(_ rect: CGRect) {
        self.init(x: rect.origin.x, y: rect.origin.y, width: rect.width, height: rect.height)
    }

    var cgRect: CGRect {
        CGRect(x: x, y: y, width: width, height: height)
    }
}