ReadViewSDK/Sources/RDAIReaderView/NaturalLanguage/RDAINaturalLanguageAnalyzer.swift

import Foundation
import NaturalLanguage

public struct RDAINaturalLanguageAnalyzer: RDAIAnalyzing {
    public struct Configuration: Sendable {
        public var maximumPassageUTF16Length: Int
        public var overlapSentenceCount: Int

        public init(maximumPassageUTF16Length: Int = 800, overlapSentenceCount: Int = 1) {
            self.maximumPassageUTF16Length = max(160, maximumPassageUTF16Length)
            self.overlapSentenceCount = max(0, overlapSentenceCount)
        }
    }

    public let configuration: Configuration

    public init(configuration: Configuration = .init()) {
        self.configuration = configuration
    }

    public func analyze(document: RDAIDocumentDescriptor, snapshot: RDAIResourceSnapshot) async -> RDAIAnalysisResult {
        await Task.detached(priority: .utility) {
            let language = Self.detectLanguage(in: snapshot.sourceText)
            let passages = Self.makePassages(document: document, snapshot: snapshot, language: language, configuration: configuration)
            let mentions = Self.makeEntityMentions(document: document, snapshot: snapshot)
            return RDAIAnalysisResult(passages: passages, entityMentions: mentions)
        }.value
    }

    private static func detectLanguage(in text: String) -> String? {
        let recognizer = NLLanguageRecognizer()
        recognizer.processString(text)
        return recognizer.dominantLanguage?.rawValue
    }

    private static func makePassages(
        document: RDAIDocumentDescriptor,
        snapshot: RDAIResourceSnapshot,
        language: String?,
        configuration: Configuration
    ) -> [RDAIPassage] {
        let sentenceRanges = sentenceRanges(
            in: snapshot.sourceText,
            maximumUTF16Length: configuration.maximumPassageUTF16Length
        )
        guard !sentenceRanges.isEmpty else { return [] }
        var result: [RDAIPassage] = []
        var startIndex = 0
        var order = 0

        while startIndex < sentenceRanges.count {
            var endIndex = startIndex
            var length = 0
            while endIndex < sentenceRanges.count {
                let candidate = sentenceRanges[endIndex]
                let nextLength = max(candidate.upperBound - sentenceRanges[startIndex].location, candidate.length)
                if endIndex > startIndex && nextLength > configuration.maximumPassageUTF16Length { break }
                length = nextLength
                endIndex += 1
            }
            guard length > 0 else { break }
            let rawRange = RDAITextRange(location: sentenceRanges[startIndex].location, length: length)
            guard let range = trimmedRange(rawRange, in: snapshot.sourceText),
                  let locator = RDAILocatorBuilder.makeLocator(document: document, snapshot: snapshot, range: range) else {
                startIndex = max(startIndex + 1, endIndex)
                continue
            }
            let text = (snapshot.sourceText as NSString).substring(with: NSRange(location: range.location, length: range.length))
            let contentHash = RDAIContentHasher.hash(text)
            result.append(RDAIPassage(
                id: RDAIContentHasher.hash("\(document.identifier.rawValue)|\(snapshot.descriptor.identifier.rawValue)|\(range.location)|\(range.length)|\(contentHash)"),
                documentIdentifier: document.identifier,
                resourceIdentifier: snapshot.descriptor.identifier,
                text: text,
                languageCode: language,
                locator: locator,
                contentHash: contentHash,
                order: order
            ))
            order += 1
            let nextStart = max(endIndex - configuration.overlapSentenceCount, startIndex + 1)
            startIndex = nextStart
        }
        return result
    }

    private static func sentenceRanges(in text: String, maximumUTF16Length: Int) -> [RDAITextRange] {
        let tokenizer = NLTokenizer(unit: .sentence)
        tokenizer.string = text
        var ranges: [RDAITextRange] = []
        tokenizer.enumerateTokens(in: text.startIndex..<text.endIndex) { range, _ in
            let nsRange = NSRange(range, in: text)
            if nsRange.length > 0 { ranges.append(RDAITextRange(location: nsRange.location, length: nsRange.length)) }
            return true
        }
        if ranges.isEmpty, !text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
            ranges = [RDAITextRange(location: 0, length: text.utf16.count)]
        }
        return ranges.flatMap { split($0, in: text, maximumUTF16Length: maximumUTF16Length) }
    }

    private static func split(
        _ range: RDAITextRange,
        in text: String,
        maximumUTF16Length: Int
    ) -> [RDAITextRange] {
        guard range.length > maximumUTF16Length else { return [range] }
        let source = text as NSString
        var result: [RDAITextRange] = []
        var cursor = range.location
        while cursor < range.upperBound {
            let candidate = min(cursor + maximumUTF16Length, range.upperBound)
            var boundary = candidate
            if candidate < range.upperBound {
                let composed = source.rangeOfComposedCharacterSequence(at: candidate)
                boundary = composed.location > cursor ? composed.location : min(composed.upperBound, range.upperBound)
            }
            guard boundary > cursor else { break }
            result.append(RDAITextRange(location: cursor, length: boundary - cursor))
            cursor = boundary
        }
        return result
    }

    private static func trimmedRange(_ range: RDAITextRange, in text: String) -> RDAITextRange? {
        let source = (text as NSString).substring(with: NSRange(location: range.location, length: range.length))
        let trimmed = source.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty else { return nil }
        let leading = source.utf16.count - source.drop(while: { $0.isWhitespace || $0.isNewline }).utf16.count
        return RDAITextRange(location: range.location + leading, length: trimmed.utf16.count)
    }

    private static func makeEntityMentions(
        document: RDAIDocumentDescriptor,
        snapshot: RDAIResourceSnapshot
    ) -> [RDAIEntityMention] {
        let tagger = NLTagger(tagSchemes: [.nameType])
        tagger.string = snapshot.sourceText
        var mentions: [RDAIEntityMention] = []
        let fullRange = snapshot.sourceText.startIndex..<snapshot.sourceText.endIndex
        tagger.enumerateTags(in: fullRange, unit: .word, scheme: .nameType, options: [.omitWhitespace, .omitPunctuation, .joinNames]) { tag, range in
            guard let tag,
                  let kind = entityKind(for: tag) else { return true }
            let nsRange = NSRange(range, in: snapshot.sourceText)
            let textRange = RDAITextRange(location: nsRange.location, length: nsRange.length)
            guard let locator = RDAILocatorBuilder.makeLocator(document: document, snapshot: snapshot, range: textRange) else { return true }
            let surfaceText = String(snapshot.sourceText[range])
            let normalized = surfaceText.trimmingCharacters(in: .whitespacesAndNewlines)
            guard !normalized.isEmpty else { return true }
            mentions.append(RDAIEntityMention(
                id: RDAIContentHasher.hash("\(snapshot.descriptor.identifier.rawValue)|\(textRange.location)|\(normalized)|\(kind.rawValue)"),
                normalizedName: normalized.folding(options: [.caseInsensitive, .diacriticInsensitive], locale: .current),
                surfaceText: normalized,
                kind: kind,
                confidence: 0.8,
                locator: locator
            ))
            return true
        }
        return mentions
    }

    private static func entityKind(for tag: NLTag) -> RDAIEntityKind? {
        switch tag {
        case .personalName: return .person
        case .placeName: return .place
        case .organizationName: return .organization
        default: return nil
        }
    }
}