//
//  Chunker.swift
//  Zyquo Atlas
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Splits a PageContext's markdown into heading-aware chunks for long-page
//  summarization (map-reduce) and chat-with-page retrieval (docs/
//  AI-BROWSER-RESEARCH.md §3). Each chunk carries its heading path and
//  character range so answers can cite and scroll-to-highlight the source.
//

import Foundation

struct Chunk: Identifiable, Hashable {
    let id: Int
    /// e.g. "Setup > macOS".
    let headingPath: String
    let text: String
    let charRange: Range<Int>
    var estimatedTokens: Int { text.count / 4 }
}

enum Chunker {
    /// ~600-token Q&A chunks by default; pass a larger target for summarization
    /// (map-reduce tolerates coarse chunks).
    static func chunk(_ context: PageContext, targetTokens: Int = 600, overlapTokens: Int = 50) -> [Chunk] {
        let md = context.markdown
        guard !md.isEmpty else { return [] }

        let targetChars = targetTokens * 4
        let overlapChars = overlapTokens * 4
        let scalars = Array(md)

        // Section boundaries from the heading index (fall back to the whole doc).
        var boundaries = context.headings.map { $0.offset }.sorted()
        if boundaries.first != 0 { boundaries.insert(0, at: 0) }
        boundaries.append(scalars.count)

        var chunks: [Chunk] = []
        var id = 0
        var headingStack: [(level: Int, title: String)] = []

        for i in 0..<(boundaries.count - 1) {
            let start = boundaries[i]
            let end = boundaries[i + 1]
            guard start < end else { continue }

            // Track heading path from the heading that opens this section.
            if let h = context.headings.first(where: { $0.offset == start }) {
                while let last = headingStack.last, last.level >= h.level { headingStack.removeLast() }
                headingStack.append((h.level, h.title))
            }
            let path = headingStack.map(\.title).joined(separator: " > ")

            // Sub-split large sections with overlap.
            var cursor = start
            while cursor < end {
                let sliceEnd = min(cursor + targetChars, end)
                let text = String(scalars[cursor..<sliceEnd]).trimmingCharacters(in: .whitespacesAndNewlines)
                if !text.isEmpty {
                    chunks.append(Chunk(id: id, headingPath: path, text: text,
                                        charRange: cursor..<sliceEnd))
                    id += 1
                }
                if sliceEnd >= end { break }
                cursor = max(cursor + targetChars - overlapChars, cursor + 1)
            }
        }
        return chunks
    }
}