# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Callable

from .parsers import ParsedPage

# Match "Figure 1:", "Figure 1.2:", "Fig. 3.", "Table 4:" etc. at line-start,
# tolerating leading bold markers. Breaks chunks BEFORE such captions so the
# caption starts its own chunk — dense embeddings pool over the whole chunk, so
# figure refs buried at the end get diluted by surrounding body text.
_FIGURE_BOUNDARY_RE = re.compile(
    # Number forms: "1", "12", "1.2", "B.1" (appendix); bold wrappers tolerated.
    r"^\**(?:Figure|Fig\.|Table|Tab\.)\s+[A-Z]?\.?\d+(?:\.\d+)?\**[\.:]",
    re.MULTILINE | re.IGNORECASE,
)


def _split_at_figure_boundaries(text: str) -> list[str]:
    """Split markdown at the start of each figure / table caption.

    Each segment starts with either the original text head or a "Figure N:" /
    "Table N:" line, so the caption anchors the embedding of the chunk it
    lands in. Returns the original text as a single-element list when no
    captions are found.
    """
    matches = list(_FIGURE_BOUNDARY_RE.finditer(text))
    if not matches:
        return [text]
    segments: list[str] = []
    last = 0
    for m in matches:
        if m.start() > last:
            segments.append(text[last : m.start()])
        last = m.start()
    segments.append(text[last:])
    return [s for s in segments if s.strip()]


@dataclass(frozen = True)
class Chunk:
    text: str
    token_count: int
    page_number: int | None = None
    source_page_index: int | None = None
    page_char_start: int | None = None
    page_char_end: int | None = None
    line_start: int | None = None
    line_end: int | None = None


TokenCounter = Callable[[str], int]


def _char_token_estimate(text: str) -> int:
    return max(1, (len(text) + 3) // 4)


def _split_on(text: str, separator: str) -> list[str]:
    if separator == "":
        return list(text)
    parts = text.split(separator)
    if len(parts) == 1:
        return parts
    glued: list[str] = []
    for i, part in enumerate(parts):
        if i < len(parts) - 1:
            glued.append(part + separator)
        else:
            if part:
                glued.append(part)
    return [p for p in glued if p]


def _atomic_split(
    text: str,
    separators: tuple[str, ...],
    max_tokens: int,
    count: TokenCounter,
) -> list[str]:
    if count(text) <= max_tokens:
        return [text]
    for sep in separators:
        pieces = _split_on(text, sep)
        if len(pieces) <= 1:
            continue
        out: list[str] = []
        for piece in pieces:
            if count(piece) <= max_tokens:
                out.append(piece)
            else:
                tail = separators[separators.index(sep) + 1 :]
                out.extend(_atomic_split(piece, tail, max_tokens, count))
        return out
    approx_chars = max(1, max_tokens * 4)
    return [text[i : i + approx_chars] for i in range(0, len(text), approx_chars)]


def _merge(
    pieces: list[str],
    max_tokens: int,
    overlap_tokens: int,
    count: TokenCounter,
) -> list[str]:
    """Greedy-merge into <= max_tokens chunks with overlap."""
    chunks: list[str] = []
    buffer: list[str] = []
    buffer_tokens = 0
    for piece in pieces:
        piece_tokens = count(piece)
        if buffer and buffer_tokens + piece_tokens > max_tokens:
            chunks.append("".join(buffer))
            if overlap_tokens > 0:
                overlap: list[str] = []
                running = 0
                for prev in reversed(buffer):
                    prev_tokens = count(prev)
                    if running + prev_tokens > overlap_tokens:
                        break
                    overlap.insert(0, prev)
                    running += prev_tokens
                buffer = list(overlap)
                buffer_tokens = running
            else:
                buffer = []
                buffer_tokens = 0
        buffer.append(piece)
        buffer_tokens += piece_tokens
    if buffer:
        chunks.append("".join(buffer))
    return [c.strip() for c in chunks if c.strip()]


def _line_bounds(text: str, start: int, end: int) -> tuple[int, int]:
    """Return 1-based inclusive line numbers for a page-local span."""
    line_start = text.count("\n", 0, start) + 1
    line_end = text.count("\n", 0, max(start, end - 1)) + 1
    return line_start, line_end


def _locate_piece(
    page_text: str,
    piece: str,
    search_cursor: int,
) -> tuple[int | None, int | None, int | None, int | None, int]:
    idx = page_text.find(piece, search_cursor)
    if idx < 0:
        idx = page_text.find(piece)
    if idx < 0:
        return None, None, None, None, search_cursor
    end = idx + len(piece)
    line_start, line_end = _line_bounds(page_text, idx, end)
    return idx, end, line_start, line_end, idx + 1


# Markdown headings first so layout-aware parser output splits at sections.
DEFAULT_SEPARATORS: tuple[str, ...] = (
    "\n# ",
    "\n## ",
    "\n### ",
    "\n#### ",
    "\n\n",
    "\n",
    ". ",
    " ",
    "",
)


def chunk_pages(
    pages: list[ParsedPage],
    *,
    max_tokens: int,
    overlap_tokens: int,
    token_counter: TokenCounter | None = None,
    separators: tuple[str, ...] = DEFAULT_SEPARATORS,
) -> list[Chunk]:
    """Split pages independently so page_number stays attached to chunks."""
    count = token_counter or _char_token_estimate
    out: list[Chunk] = []
    for page_index, page in enumerate(pages):
        search_cursor = 0
        for segment in _split_at_figure_boundaries(page.text):
            atomic = _atomic_split(segment, separators, max_tokens, count)
            merged = _merge(atomic, max_tokens, overlap_tokens, count)
            for piece in merged:
                start, end, line_start, line_end, search_cursor = _locate_piece(
                    page.text,
                    piece,
                    search_cursor,
                )
                out.append(
                    Chunk(
                        text = piece,
                        token_count = count(piece),
                        page_number = page.page_number,
                        source_page_index = page_index,
                        page_char_start = start,
                        page_char_end = end,
                        line_start = line_start,
                        line_end = line_end,
                    )
                )
    return out


_PAGE_SEPARATOR = "\n\n"


def chunk_pages_with_spans(
    pages: list[ParsedPage],
    *,
    max_tokens: int,
    overlap_tokens: int,
    token_counter: TokenCounter | None = None,
    separators: tuple[str, ...] = DEFAULT_SEPARATORS,
) -> tuple[str, list[Chunk], list[tuple[int, int]]]:
    """Late-chunking variant: joins pages so the embedder sees the whole doc.

    Returns ``(full_doc, chunks, char_spans)``; ``char_spans[i]`` is the
    (start, end) char offset of ``chunks[i].text`` inside ``full_doc``.
    Page numbers are recovered by overlap with the original page ranges.
    """
    count = token_counter or _char_token_estimate

    parts: list[str] = []
    page_ranges: list[tuple[int, int, int, int | None]] = []
    cursor = 0
    for index, page in enumerate(pages):
        parts.append(page.text)
        start = cursor
        end = cursor + len(page.text)
        page_ranges.append((start, end, index, page.page_number))
        cursor = end
        if index < len(pages) - 1:
            cursor += len(_PAGE_SEPARATOR)
    full_doc = _PAGE_SEPARATOR.join(parts)

    atomic: list[str] = []
    for segment in _split_at_figure_boundaries(full_doc):
        atomic.extend(_atomic_split(segment, separators, max_tokens, count))
    merged = _merge(atomic, max_tokens, overlap_tokens, count)

    chunks: list[Chunk] = []
    char_spans: list[tuple[int, int]] = []
    search_cursor = 0
    for piece in merged:
        text = piece.strip()
        if not text:
            continue
        idx = full_doc.find(text, search_cursor)
        if idx < 0:
            # Overlap can push past a chunk's true start; restart from head.
            idx = full_doc.find(text)
        if idx < 0:
            continue
        end_idx = idx + len(text)
        page_locator = _page_for_span(idx, end_idx, page_ranges)
        source_page_index: int | None = None
        page_number: int | None = None
        page_char_start: int | None = None
        page_char_end: int | None = None
        line_start: int | None = None
        line_end: int | None = None
        if page_locator is not None:
            page_start, page_end, page_idx, page_no = page_locator
            source_page_index = page_idx
            page_number = page_no
            page_char_start = max(0, idx - page_start)
            page_char_end = min(page_end, end_idx) - page_start
            line_start, line_end = _line_bounds(
                pages[page_idx].text,
                page_char_start,
                page_char_end,
            )
        chunks.append(
            Chunk(
                text = text,
                token_count = count(text),
                page_number = page_number,
                source_page_index = source_page_index,
                page_char_start = page_char_start,
                page_char_end = page_char_end,
                line_start = line_start,
                line_end = line_end,
            )
        )
        char_spans.append((idx, end_idx))
        # Advance past start (not end) so overlapping next chunk is findable.
        search_cursor = idx + 1

    return full_doc, chunks, char_spans


def _page_for_span(
    start: int,
    end: int,
    page_ranges: list[tuple[int, int, int, int | None]],
) -> tuple[int, int, int, int | None] | None:
    for ps, pe, page_index, page_number in page_ranges:
        if start < pe and end > ps:
            return ps, pe, page_index, page_number
    return None