"""Unit tests for src.extraction_v2.mermaid_chunking.

Covers the I/O matrix entries from
docs/implementation-artifacts/spec-textiq-625-unified-mermaid-aware-chunking.md:

  - chunk_with_atomic_mermaid: plain prose, mermaid in-context, oversize fence
    split, oversize single line → ValueError, tail-merge of <MIN_CHUNK_TOKENS,
    cross-section merge never happens
  - split_oversize_mermaid: balanced re-wrap of every part
  - prune_header_only_nodes: bare ATX headers dropped, ancestor cascade
  - build_text_parser_map / build_excel_parser_map: shape sanity
"""

from __future__ import annotations

import sys
from pathlib import Path

import pytest

REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT))

from src.extraction_v2.mermaid_chunking import (
    EXCEL_PARSER_IDS,
    FENCE_END_RE,
    HEADER_ONLY_RE,
    MERMAID_FENCE_START_RE,
    MIN_CHUNK_TOKENS,
    TEXT_PARSER_IDS,
    build_excel_parser_map,
    build_text_parser_map,
    chunk_with_atomic_mermaid,
    prune_header_only_nodes,
    split_oversize_mermaid,
)


def _word_tokenize(text: str) -> int:
    """Cheap word-count tokenizer so tests don't depend on tiktoken at runtime."""
    return max(1, len(text.split()))


def _word_tokenize_list(text: str):
    return text.split()


class TestRegexes:
    def test_mermaid_fence_start_matches_with_attrs(self):
        assert MERMAID_FENCE_START_RE.match("```mermaid")
        assert MERMAID_FENCE_START_RE.match("```mermaid theme=dark")
        assert not MERMAID_FENCE_START_RE.match("```python")
        assert not MERMAID_FENCE_START_RE.match("  ```mermaid")  # must start at col 0

    def test_fence_end_strict(self):
        assert FENCE_END_RE.match("```")
        assert FENCE_END_RE.match("```   ")
        assert not FENCE_END_RE.match("``` mermaid")
        assert not FENCE_END_RE.match("text ```")

    def test_header_only_re(self):
        assert HEADER_ONLY_RE.match("# top")
        assert HEADER_ONLY_RE.match("###### deep")
        assert HEADER_ONLY_RE.match("  ## indented header")
        assert not HEADER_ONLY_RE.match("## header\nthen body")
        assert not HEADER_ONLY_RE.match("plain text")
        assert not HEADER_ONLY_RE.match("####### too-deep")  # > 6 hashes


class TestChunkWithAtomicMermaidPlainProse:
    def test_short_prose_returns_single_chunk(self):
        text = "alpha\nbeta\ngamma\n"
        out = chunk_with_atomic_mermaid(text, target=100, max_size=200, tokenize=_word_tokenize)
        assert len(out) == 1
        chunk_text, has_mermaid = out[0]
        assert "alpha" in chunk_text and "gamma" in chunk_text
        assert has_mermaid is False

    def test_prose_splits_when_over_target(self):
        # Each line is 1 word; with target=2, expect ~3 chunks. Use bigger
        # MIN_CHUNK_TOKENS bypass by making lines weigh > MIN_CHUNK_TOKENS each.
        # Easier: target=2 with 6 single-word lines -> 3 chunks pre-merge,
        # post-merge merges trailing 1-line tail back if it's < MIN.
        text = "alpha\nbeta\ngamma\ndelta\nepsilon\nzeta\n"
        out = chunk_with_atomic_mermaid(
            text, target=2, max_size=200, tokenize=_word_tokenize
        )
        # All chunks must contain only prose lines from the input.
        for c, has_m in out:
            assert has_m is False
            assert all(line in {"alpha", "beta", "gamma", "delta", "epsilon", "zeta", ""} for line in c.split("\n"))
        assert len(out) >= 1


class TestChunkWithAtomicMermaidInContext:
    def test_fence_kept_in_same_chunk_as_prose(self):
        text = (
            "intro line one\n"
            "intro line two\n"
            "```mermaid\n"
            "flowchart TD\n"
            "  a --> b\n"
            "```\n"
            "trailing line\n"
        )
        out = chunk_with_atomic_mermaid(
            text, target=100, max_size=200, tokenize=_word_tokenize
        )
        assert len(out) == 1
        chunk_text, has_mermaid = out[0]
        assert has_mermaid is True
        assert "```mermaid" in chunk_text
        assert chunk_text.count("```") == 2  # opener + closer, balanced

    def test_fence_chunk_can_exceed_target_up_to_max(self):
        # target=1 but max=200 -> with a fence, chunker keeps it whole.
        # We pad prose to ensure cur_tok > target before the fence arrives.
        text = (
            "prose one\n"
            "prose two\n"
            "```mermaid\n"
            "flowchart TD\n"
            "  x --> y\n"
            "```\n"
        )
        out = chunk_with_atomic_mermaid(
            text, target=1, max_size=200, tokenize=_word_tokenize
        )
        # We expect at most 2 chunks (prose flushed before fence, then fence
        # in its own chunk OR the fence absorbed the prose). Either way every
        # fence-bearing chunk has balanced ```'s.
        for c, has_m in out:
            if has_m:
                assert c.count("```") == 2

    def test_mermaid_only_input(self):
        text = "```mermaid\nflowchart TD\n  a --> b\n```\n"
        out = chunk_with_atomic_mermaid(
            text, target=10, max_size=100, tokenize=_word_tokenize
        )
        assert len(out) == 1
        chunk_text, has_mermaid = out[0]
        assert has_mermaid is True
        assert chunk_text.startswith("```mermaid")
        assert "```" in chunk_text.splitlines()[-1] or chunk_text.rstrip().endswith("```")


class TestOversizeFence:
    def test_oversize_fence_split_internally(self):
        # 10 body lines * 2 words = 20 body tokens; with max=8 and wrapper~3
        # tokens, the chunker should split into multiple ```mermaid wrappers.
        body = "\n".join(f"node{i} word" for i in range(10))
        text = f"```mermaid\n{body}\n```\n"
        out = chunk_with_atomic_mermaid(
            text, target=4, max_size=8, tokenize=_word_tokenize
        )
        # Every part of the split must be balanced and labelled has_mermaid.
        assert len(out) >= 2
        for chunk_text, has_mermaid in out:
            assert has_mermaid is True
            assert chunk_text.count("```") == 2

    def test_oversize_single_line_raises(self):
        # Wrapper alone ~3 tokens; body has one line worth 20 tokens; limit=5.
        body_line = "word " * 20
        block_lines = ["```mermaid", body_line.strip(), "```"]
        with pytest.raises(ValueError, match="single mermaid line exceeds embedding context"):
            split_oversize_mermaid(block_lines, _word_tokenize, limit=5)

    def test_split_oversize_mermaid_preserves_opener_attrs(self):
        block_lines = [
            "```mermaid theme=dark",
            "flowchart TD",
            "  a --> b",
            "  c --> d",
            "```",
        ]
        parts = split_oversize_mermaid(block_lines, _word_tokenize, limit=6)
        assert len(parts) >= 2
        for p in parts:
            assert p.startswith("```mermaid theme=dark\n")
            assert p.rstrip().endswith("```")


class TestTailMerge:
    def test_small_tail_merges_into_previous(self):
        # Two-line prose with target=1: first chunk = "alpha", tail = "a"
        # (1 token, < MIN_CHUNK_TOKENS). Tail should merge backward.
        text = "alpha word word word word\nb\n"  # second line is 1 tok, < MIN=64
        out = chunk_with_atomic_mermaid(
            text, target=2, max_size=200, tokenize=_word_tokenize
        )
        # After merge expected: 1 chunk containing both lines.
        assert len(out) == 1
        joined, _ = out[0]
        assert "alpha" in joined and "\nb" in joined

    def test_no_merge_when_combined_exceeds_max(self):
        # First chunk near max; second chunk small but combined > max.
        # Build first chunk to be ~max_size, second to be < MIN.
        big = "w " * 100  # 100 tokens (single line)
        text = f"{big.strip()}\nshort\n"
        out = chunk_with_atomic_mermaid(
            text, target=99, max_size=100, tokenize=_word_tokenize
        )
        # Combined would be 101 > max=100, so should not merge.
        assert len(out) == 2


class TestPruneHeaderOnly:
    def test_drops_leaf_header_only_node(self):
        from llama_index.core.schema import TextNode

        nodes = [
            TextNode(text="## just a header"),
            TextNode(text="this is real prose content"),
        ]
        kept = prune_header_only_nodes(nodes)
        assert len(kept) == 1
        assert kept[0].text == "this is real prose content"

    def test_keeps_section_with_body(self):
        from llama_index.core.schema import TextNode

        nodes = [
            TextNode(text="## Section A\n\nThe body of the section."),
        ]
        kept = prune_header_only_nodes(nodes)
        assert len(kept) == 1

    def test_cascading_parent_drop(self):
        from llama_index.core.schema import (
            NodeRelationship,
            RelatedNodeInfo,
            TextNode,
        )

        parent = TextNode(text="## Parent header")
        child = TextNode(text="### Child header only")
        parent.relationships[NodeRelationship.CHILD] = [
            RelatedNodeInfo(node_id=child.node_id)
        ]
        child.relationships[NodeRelationship.PARENT] = RelatedNodeInfo(
            node_id=parent.node_id
        )
        kept = prune_header_only_nodes([parent, child])
        assert kept == []  # parent dropped because its only child was dropped


class TestParserMaps:
    def test_text_parser_map_has_two_levels(self):
        # We supply a fake tokenizer; SentenceSplitter only stores it.
        m = build_text_parser_map(_word_tokenize_list)
        assert set(m.keys()) == {"markdown", "chunk_size_512"}
        from llama_index.core.node_parser import MarkdownNodeParser, SentenceSplitter
        assert isinstance(m["markdown"], MarkdownNodeParser)
        assert isinstance(m["chunk_size_512"], SentenceSplitter)
        assert m["chunk_size_512"].paragraph_separator == "\n\n"
        # 20% overlap as configured in spec.
        assert m["chunk_size_512"].chunk_overlap == 512 // 5

    def test_excel_parser_map_keys_match_parser_ids(self):
        m = build_excel_parser_map(_word_tokenize, embed_max_tokens=8191)
        assert list(m.keys()) == EXCEL_PARSER_IDS
        from llama_index.core.node_parser import MarkdownNodeParser
        assert isinstance(m["markdown"], MarkdownNodeParser)
        # The splitter exposes target_tokens=512 from spec.
        splitter = m["chunk_size_512"]
        assert splitter.target_tokens == 512
        assert splitter.max_tokens == 8191

    def test_parser_id_constants(self):
        assert EXCEL_PARSER_IDS == ["markdown", "chunk_size_512"]
        assert TEXT_PARSER_IDS == ["markdown", "chunk_size_512"]


class TestEndToEndExcelHierarchy:
    def test_hierarchical_parser_with_excel_map_strips_header_only_section(self):
        from llama_index.core import Document
        from llama_index.core.node_parser import HierarchicalNodeParser, get_leaf_nodes

        md = (
            "# Sheet1\n\n"
            "## Section A\n\n"
            "Real prose line one.\n"
            "Real prose line two.\n\n"
            "## Section B\n\n"
            "```mermaid\nflowchart TD\n  a --> b\n```\n\n"
            "Closing prose.\n\n"
            "## Empty Header\n"
        )
        parser_map = build_excel_parser_map(_word_tokenize, embed_max_tokens=8191)
        hp = HierarchicalNodeParser(
            node_parser_ids=EXCEL_PARSER_IDS, node_parser_map=parser_map
        )
        nodes = hp.get_nodes_from_documents([Document(text=md, metadata={"sheet_name": "Sheet1"})])
        nodes = list(prune_header_only_nodes(nodes))
        leaves = get_leaf_nodes(nodes)

        # Empty Header section (and its leaf) should be pruned.
        assert all("Empty Header" not in n.text for n in nodes)

        # Sheet metadata flows through.
        assert all(n.metadata.get("sheet_name") == "Sheet1" for n in leaves)

        # At least one mermaid-bearing leaf.
        mermaid_leaves = [n for n in leaves if "```mermaid" in n.text]
        assert mermaid_leaves
        for n in mermaid_leaves:
            assert n.text.count("```") == 2  # balanced
            assert n.metadata.get("has_mermaid") is True