"""Mermaid-aware chunking helpers shared by hierarchical_chunk_* DAG tasks.

The line-based custom chunker (`chunk_with_atomic_mermaid`) is used **only**
by the Excel branch: Docling-xlsx markdown has almost no `\n\n` paragraph
anchors, so the stock `SentenceSplitter(paragraph_separator="\n\n")` falls
through to char-level splits that bisect table rows mid-cell. Word / PDF /
PPTX and `.md` inputs have normal paragraph structure and use a stock
2-level `MarkdownNodeParser -> SentenceSplitter(512)` hierarchy via
`build_text_parser_map`.

Public API:
    chunk_with_atomic_mermaid(text, target, max_size, tokenize)
        Excel-only line-based chunker. Treats each ```mermaid fence as an
        atomic line-group; pure-prose chunks stay <= target, chunks that
        contain a fence may grow up to max_size. Backward-merges tail
        chunks < MIN_CHUNK_TOKENS within the same section.
    split_oversize_mermaid(block_lines, tokenize, limit)
        Split a single fence whose total tokens exceed `limit` along its
        body lines, rewrapping each part as its own ```mermaid fence.
    prune_header_only_nodes(nodes)
        Drop any node whose stripped text matches HEADER_ONLY_RE; remove
        ancestors whose entire subtree was pruned.
    build_excel_parser_map(tokenize_fn, embed_max_tokens, target_tokens=512)
        Parser map for the Excel branch ({markdown, chunk_size_512}).
    build_text_parser_map(tokenizer)
        Stock 2-level parser map for non-md / md branches
        ({markdown, chunk_size_512}).
"""

from __future__ import annotations

import re
from typing import Any, Callable, Dict, List, Sequence, Tuple

MERMAID_FENCE_START_RE = re.compile(r"^```mermaid\b")
FENCE_END_RE = re.compile(r"^```\s*$")
HEADER_ONLY_RE = re.compile(r"^\s*#{1,6} [^\n]+\s*$")

MIN_CHUNK_TOKENS = 64

EXCEL_PARSER_IDS: List[str] = ["markdown", "chunk_size_512"]
TEXT_PARSER_IDS: List[str] = ["markdown", "chunk_size_512"]


def split_oversize_mermaid(
    block_lines: List[str], tokenize: Callable[[str], int], limit: int
) -> List[str]:
    """Split a single mermaid fence whose total tokens exceed `limit`.

    Splits along the block's inner body lines and rewraps each part as its
    own ```mermaid...``` fence. The opener (e.g. ```` ```mermaid ````) and
    closer (```` ``` ````) are preserved verbatim per part.

    Raises ValueError if a single inner body line by itself exceeds the
    limit (i.e. cannot be made to fit even alone inside the wrapper).
    """
    opener, closer = block_lines[0], block_lines[-1]
    inner = block_lines[1:-1]
    wrapper_tok = tokenize(opener + "\n\n" + closer)
    parts: List[str] = []
    cur: List[str] = []
    cur_tok = wrapper_tok
    for ln in inner:
        ln_tok = tokenize(ln + "\n")
        if wrapper_tok + ln_tok > limit:
            raise ValueError(
                f"single mermaid line exceeds embedding context ({ln_tok} tokens)"
            )
        if cur_tok + ln_tok > limit and cur:
            parts.append(opener + "\n" + "\n".join(cur) + "\n" + closer)
            cur = [ln]
            cur_tok = wrapper_tok + ln_tok
        else:
            cur.append(ln)
            cur_tok += ln_tok
    if cur:
        parts.append(opener + "\n" + "\n".join(cur) + "\n" + closer)
    return parts


def chunk_with_atomic_mermaid(
    text: str,
    target: int,
    max_size: int,
    tokenize: Callable[[str], int],
) -> List[Tuple[str, bool]]:
    """Line-based chunker that treats ```mermaid fences as atomic line-groups.

    Returns a list of `(chunk_text, contains_mermaid)` tuples. Pure-prose
    chunks stay <= `target`. A chunk containing at least one mermaid fence
    may grow up to `max_size` so the diagram stays adjacent to its
    explanatory prose. A single fence > `max_size` is split internally via
    `split_oversize_mermaid`.

    After chunking, a backward-merge pass folds chunks < MIN_CHUNK_TOKENS
    into the previous chunk if the combined size still fits within
    `max_size`. This eliminates tail-fragment leaves (e.g. lone table rows
    that fell past a flush boundary). Cross-section merging is the caller's
    responsibility -- this function never sees section boundaries.
    """
    lines = text.split("\n")
    chunks: List[Tuple[str, bool]] = []
    cur: List[str] = []
    cur_tok = 0
    cur_has_mermaid = False

    def flush() -> None:
        nonlocal cur, cur_tok, cur_has_mermaid
        if cur:
            chunks.append(("\n".join(cur), cur_has_mermaid))
        cur, cur_tok, cur_has_mermaid = [], 0, False

    i = 0
    n = len(lines)
    while i < n:
        ln = lines[i]
        if MERMAID_FENCE_START_RE.match(ln):
            # Locate the closing ``` line; tolerate unterminated fences by
            # consuming to end-of-text.
            j = i + 1
            while j < n and not FENCE_END_RE.match(lines[j]):
                j += 1
            end = j if j < n else n - 1
            block_lines = lines[i:end + 1]
            block_text = "\n".join(block_lines)
            block_tok = tokenize(block_text)

            if block_tok > max_size:
                flush()
                for part in split_oversize_mermaid(block_lines, tokenize, max_size):
                    chunks.append((part, True))
                i = end + 1
                continue

            if cur_tok + block_tok > max_size and cur:
                flush()
            cur.extend(block_lines)
            cur_tok += block_tok
            cur_has_mermaid = True
            i = end + 1
        else:
            tok_n = tokenize(ln + "\n")
            limit = max_size if cur_has_mermaid else target
            if cur_tok + tok_n > limit and cur:
                flush()
            cur.append(ln)
            cur_tok += tok_n
            i += 1
    flush()

    # Backward tail-merge within the supplied section.
    merged: List[Tuple[str, bool]] = []
    for chunk_text, has_m in chunks:
        if merged and tokenize(chunk_text) < MIN_CHUNK_TOKENS:
            prev_text, prev_has = merged[-1]
            combined = prev_text + "\n" + chunk_text
            if tokenize(combined) <= max_size:
                merged[-1] = (combined, prev_has or has_m)
                continue
        merged.append((chunk_text, has_m))
    return merged


def prune_header_only_nodes(nodes: Sequence[Any]) -> List[Any]:
    """Drop nodes whose stripped text is just an ATX header line.

    Iteratively removes nodes matching HEADER_ONLY_RE plus any ancestor
    whose entire surviving subtree was pruned. Preserves topological order.
    Operates on llama_index BaseNode-like objects (uses `node_id`, `text`,
    and `relationships[CHILD]`).
    """
    from llama_index.core.schema import NodeRelationship

    drop: set = set()
    for n in nodes:
        if HEADER_ONLY_RE.match(n.text):
            drop.add(n.node_id)

    changed = True
    while changed:
        changed = False
        for n in nodes:
            if n.node_id in drop:
                continue
            child_rel = n.relationships.get(NodeRelationship.CHILD)
            if not child_rel:
                continue
            children = child_rel if isinstance(child_rel, list) else [child_rel]
            if children and all(r.node_id in drop for r in children):
                drop.add(n.node_id)
                changed = True

    return [n for n in nodes if n.node_id not in drop]


def build_text_parser_map(tokenizer: Callable[[str], List[int]]) -> Dict[str, Any]:
    """Stock 2-level hierarchy for Word / PDF / PPTX / .md branches.

    Returns a parser map suitable for `HierarchicalNodeParser` with
    `node_parser_ids = TEXT_PARSER_IDS`. The `SentenceSplitter` uses 20%
    overlap and `paragraph_separator="\\n\\n"` (markdown anchors).
    """
    from llama_index.core.node_parser import MarkdownNodeParser, SentenceSplitter

    return {
        "markdown": MarkdownNodeParser(),
        "chunk_size_512": SentenceSplitter(
            chunk_size=512,
            chunk_overlap=int(512 * 0.2),
            tokenizer=tokenizer,
            paragraph_separator="\n\n",
        ),
    }


def build_excel_parser_map(
    tokenize_fn: Callable[[str], int],
    embed_max_tokens: int,
) -> Dict[str, Any]:
    """Parser map for the Excel branch (keys: `EXCEL_PARSER_IDS`).

    Returns a dict whose `chunk_size_512` entry is a `MermaidAwareLineSplitter`
    -- a `TextSplitter` subclass that delegates to `chunk_with_atomic_mermaid`
    and stamps `has_mermaid` metadata on every leaf. The 512-token target is
    fixed by spec; the splitter still grows mermaid-bearing chunks up to
    `embed_max_tokens`. Callers should pair this with
    `HierarchicalNodeParser(node_parser_ids=EXCEL_PARSER_IDS, ...)` and apply
    `prune_header_only_nodes` to the result.
    """
    from llama_index.core.node_parser import MarkdownNodeParser

    return {
        "markdown": MarkdownNodeParser(),
        "chunk_size_512": MermaidAwareLineSplitter(
            tokenize_fn=tokenize_fn,
            target_tokens=512,
            max_tokens=embed_max_tokens,
        ),
    }


def _make_mermaid_aware_splitter_class():
    """Defer class definition to first use so importing this module does not
    require llama_index at module-load time (helps `@task.external_python`
    environments that load the helper as a string before resolving deps).
    """
    from llama_index.core.bridge.pydantic import PrivateAttr
    from llama_index.core.node_parser import TextSplitter

    class _MermaidAwareLineSplitter(TextSplitter):
        """TextSplitter wrapping `chunk_with_atomic_mermaid` for Excel chunking.

        Field semantics:
            target_tokens: pure-prose chunk budget.
            max_tokens: hard ceiling -- mermaid-bearing chunks may grow to here.

        The `tokenize_fn` callable is held in a PrivateAttr so it bypasses
        pydantic schema validation; it is not serialized.
        """

        target_tokens: int = 512
        max_tokens: int = 8191
        _tokenize_fn: Callable[[str], int] = PrivateAttr()

        def __init__(
            self,
            *,
            tokenize_fn: Callable[[str], int],
            target_tokens: int = 512,
            max_tokens: int = 8191,
            **kwargs: Any,
        ) -> None:
            super().__init__(
                target_tokens=target_tokens, max_tokens=max_tokens, **kwargs
            )
            self._tokenize_fn = tokenize_fn

        @classmethod
        def class_name(cls) -> str:
            return "MermaidAwareLineSplitter"

        def split_text(self, text: str) -> List[str]:
            return [
                t
                for t, _ in chunk_with_atomic_mermaid(
                    text, self.target_tokens, self.max_tokens, self._tokenize_fn
                )
            ]

        def _parse_nodes(
            self,
            nodes: Sequence[Any],
            show_progress: bool = False,
            **kwargs: Any,
        ) -> List[Any]:
            out = super()._parse_nodes(nodes, show_progress=show_progress, **kwargs)
            for leaf in out:
                leaf.metadata["has_mermaid"] = "```mermaid" in leaf.text
            return out

    return _MermaidAwareLineSplitter


class _LazySplitterProxy:
    """Lazy-resolves to the real `_MermaidAwareLineSplitter` on first call.

    Avoids importing llama_index at module-load time when the helper is
    inlined into `@task.external_python` source.
    """

    _cls = None

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if _LazySplitterProxy._cls is None:
            _LazySplitterProxy._cls = _make_mermaid_aware_splitter_class()
        return _LazySplitterProxy._cls(*args, **kwargs)


MermaidAwareLineSplitter = _LazySplitterProxy()
