"""Tests for entity extraction stability after refactoring.

Covers:
- entity_types_guidance injected into prompts (text mode and JSON mode)
- custom entity_types_guidance via addon_params overrides default
- ENTITY_TYPES env var raises SystemExit at LightRAG init
- EntityExtractionResult Pydantic schema used in JSON mode (entity_extraction kwarg)
- Default entity type guidance constant is present and non-empty
"""

import json
import os
import re
from pathlib import Path
from unittest.mock import AsyncMock, patch

import pytest

from lightrag.utils import EmbeddingFunc, Tokenizer, TokenizerInterface


class DummyTokenizer(TokenizerInterface):
    """Simple 1:1 character-to-token mapping for testing."""

    def encode(self, content: str):
        return [ord(ch) for ch in content]

    def decode(self, tokens):
        return "".join(chr(token) for token in tokens)


def _make_global_config(
    addon_params: dict | None = None,
    use_json: bool = False,
    max_gleaning: int = 0,
    prompt_profile: dict | None = None,
) -> dict:
    tokenizer = Tokenizer("dummy", DummyTokenizer())
    extract_func = AsyncMock(return_value="")
    return {
        "llm_model_func": extract_func,
        "role_llm_funcs": {
            "extract": extract_func,
            "keyword": extract_func,
            "query": extract_func,
            "vlm": extract_func,
        },
        "entity_extract_max_gleaning": max_gleaning,
        "entity_extract_max_records": 100,
        "entity_extract_max_entities": 40,
        "addon_params": addon_params if addon_params is not None else {},
        "tokenizer": tokenizer,
        "max_extract_input_tokens": 20480,
        "llm_model_max_async": 1,
        "entity_extraction_use_json": use_json,
        "_entity_extraction_prompt_profile": prompt_profile,
    }


def _make_chunks(content: str = "Alice founded Acme Corp in 1990.") -> dict[str, dict]:
    return {
        "chunk-001": {
            "tokens": len(content),
            "content": content,
            "full_doc_id": "doc-001",
            "chunk_order_index": 0,
        }
    }


def _require_yaml() -> None:
    pytest.importorskip("yaml")


def _write_prompt_profile(
    path: Path,
    *,
    guidance: str | None = None,
    text_examples: list[str] | None = None,
    json_examples: list[str] | None = None,
) -> None:
    lines: list[str] = []

    def _append_block(key: str, value: str) -> None:
        lines.append(f"{key}: |")
        for line in value.strip("\n").splitlines():
            lines.append(f"  {line}")

    def _append_examples(key: str, values: list[str]) -> None:
        lines.append(f"{key}:")
        for value in values:
            lines.append("  - |")
            for line in value.strip("\n").splitlines():
                lines.append(f"    {line}")

    if guidance is not None:
        _append_block("entity_types_guidance", guidance)
    if text_examples is not None:
        _append_examples("entity_extraction_examples", text_examples)
    if json_examples is not None:
        _append_examples("entity_extraction_json_examples", json_examples)

    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def _dummy_embedding_func() -> EmbeddingFunc:
    async def _embed(texts):
        return [[0.0, 0.0, 0.0] for _ in texts]

    return EmbeddingFunc(embedding_dim=3, func=_embed)


def _patch_prompt_dir(path: Path):
    return patch("lightrag.prompt.get_entity_type_prompt_dir", return_value=path)


def _text_profile_example(label: str) -> str:
    return f"""---Entity Types---
- ExampleType: Test type

---Input Text---
```
{label}
```

---Output---
entity{{tuple_delimiter}}{label}{{tuple_delimiter}}ExampleType{{tuple_delimiter}}{label} description.
{{completion_delimiter}}"""


def _json_profile_example(label: str) -> str:
    return f"""---Entity Types---
- ExampleType: Test type

---Input Text---
```
{label}
```

---Output---
{{
  "entities": [
    {{"name": "{label}", "type": "ExampleType", "description": "{label} description."}}
  ],
  "relationships": []
}}"""


# --- Minimal valid LLM responses ---

_TEXT_MODE_RESPONSE = (
    "entity<|#|>Alice<|#|>Person<|#|>Alice is the founder of Acme Corp."
    "\nentity<|#|>Acme Corp<|#|>Organization<|#|>Acme Corp is a company founded by Alice."
    "\nrelation<|#|>Alice<|#|>Acme Corp<|#|>founded<|#|>Alice founded Acme Corp."
    "\n<|COMPLETE|>"
)

_TEXT_MODE_MISPREFIXED_RELATION_RESPONSE = (
    "entity<|#|>Alice<|#|>Person<|#|>Alice is the founder of Acme Corp."
    "\nentity<|#|>Acme Corp<|#|>Organization<|#|>Acme Corp is a company founded by Alice."
    "\nentity<|#|>Alice<|#|>Acme Corp<|#|>founded<|#|>Alice founded Acme Corp."
    "\n<|COMPLETE|>"
)

_TEXT_MODE_GLEANED_RELATION_RESPONSES = [
    _TEXT_MODE_MISPREFIXED_RELATION_RESPONSE,
    "\nrelation<|#|>Alice<|#|>Acme Corp<|#|>founded<|#|>Alice founded Acme Corp.\n<|COMPLETE|>",
]

_TEXT_MODE_CROSS_PASS_RELATION_RESPONSES = [
    "entity<|#|>Alice<|#|>Person<|#|>Alice founded a company.\n<|COMPLETE|>",
    "entity<|#|>Acme Corp<|#|>Organization<|#|>Acme Corp was founded by Alice."
    "\nrelation<|#|>Alice<|#|>Acme Corp<|#|>founded<|#|>Alice founded Acme Corp.\n<|COMPLETE|>",
]

_JSON_MODE_RESPONSE = json.dumps(
    {
        "entities": [
            {
                "name": "Alice",
                "type": "Person",
                "description": "Alice is the founder of Acme Corp.",
            },
            {
                "name": "Acme Corp",
                "type": "Organization",
                "description": "Acme Corp is a company founded by Alice.",
            },
        ],
        "relationships": [
            {
                "source": "Alice",
                "target": "Acme Corp",
                "keywords": "founded",
                "description": "Alice founded Acme Corp.",
            },
        ],
    }
)


class _DummyTextChunksStorage:
    async def get_by_id(self, chunk_id: str):
        return {"file_path": "test.md"}


# ---------------------------------------------------------------------------
# 1. Default entity_types_guidance constant
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_default_entity_types_guidance_exists():
    """PROMPTS['default_entity_types_guidance'] must be a non-empty string."""
    from lightrag.prompt import PROMPTS

    guidance = PROMPTS["default_entity_types_guidance"]
    assert isinstance(guidance, str)
    assert len(guidance.strip()) > 0


@pytest.mark.offline
def test_default_entity_types_guidance_covers_all_types():
    """Default guidance must mention all 11 canonical entity types."""
    from lightrag.prompt import PROMPTS

    guidance = PROMPTS["default_entity_types_guidance"]
    expected_types = [
        "Person",
        "Creature",
        "Organization",
        "Location",
        "Event",
        "Concept",
        "Method",
        "Content",
        "Data",
        "Artifact",
        "NaturalObject",
    ]
    for t in expected_types:
        assert t in guidance, (
            f"Expected entity type '{t}' missing from default_entity_types_guidance"
        )


@pytest.mark.offline
def test_builtin_entity_extraction_examples_are_format_only():
    """Built-in examples must be placeholder templates, not extractable demos.

    Rather than denylisting specific sample names (brittle: any new concrete
    content with different names would slip through), assert the structural
    shape of a format-only template: no per-example section headers that would
    reintroduce a sample ``---Input Text---`` / ``---Output---`` demo, and every
    data value is an angle-bracket placeholder rather than concrete prose.
    """
    from lightrag.prompt import PROMPTS

    section_markers = ("---Input Text---", "---Output---", "---Entity Types---")
    placeholder = re.compile(r"<[^<>]+>")
    tuple_delimiter = PROMPTS["DEFAULT_TUPLE_DELIMITER"]
    completion_delimiter = PROMPTS["DEFAULT_COMPLETION_DELIMITER"]

    # Text examples: every field after the leading entity/relation tag must be a
    # bare placeholder; concrete sample values would not match.
    for example in PROMPTS["entity_extraction_examples"]:
        for marker in section_markers:
            assert marker not in example
        rendered = example.format(
            tuple_delimiter=tuple_delimiter,
            completion_delimiter=completion_delimiter,
        )
        for line in rendered.splitlines():
            line = line.strip()
            if not line or line == completion_delimiter:
                continue
            tag, *fields = line.split(tuple_delimiter)
            assert tag in {"entity", "relation"}
            assert fields  # data rows must carry at least one value field
            for field in fields:
                assert placeholder.fullmatch(field), field

    # JSON examples: every entity/relationship field value must be a placeholder.
    for example in PROMPTS["entity_extraction_json_examples"]:
        for marker in section_markers:
            assert marker not in example
        parsed = json.loads(example)
        records = parsed["entities"] + parsed["relationships"]
        assert records
        for record in records:
            for value in record.values():
                assert placeholder.fullmatch(value), value


@pytest.mark.offline
def test_entity_extraction_system_prompts_label_examples_as_format_templates():
    from lightrag.prompt import PROMPTS

    for prompt_key in (
        "entity_extraction_system_prompt",
        "entity_extraction_json_system_prompt",
    ):
        prompt = PROMPTS[prompt_key]
        assert "---Output Format Template---" in prompt
        assert "---Examples---" not in prompt
        assert "output format template only" in prompt
        assert "not source text" in prompt
        assert "must never be used as extraction content" in prompt


@pytest.mark.offline
def test_text_examples_render_tuple_and_completion_delimiters():
    from lightrag.prompt import PROMPTS

    rendered = "\n".join(PROMPTS["entity_extraction_examples"]).format(
        tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
        completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
    )

    assert (
        "entity<|#|><entity_name><|#|><entity_type><|#|><entity_description>"
        in rendered
    )
    assert (
        "relation<|#|><source_entity><|#|><target_entity><|#|>"
        "<relationship_keywords><|#|><relationship_description>" in rendered
    )
    assert "<|COMPLETE|>" in rendered
    assert "{tuple_delimiter}" not in rendered
    assert "{completion_delimiter}" not in rendered


@pytest.mark.offline
def test_json_examples_are_parseable_format_templates():
    """JSON examples must be raw JSON templates with valid endpoint references."""
    from lightrag.prompt import PROMPTS

    for example in PROMPTS["entity_extraction_json_examples"]:
        parsed = json.loads(example)
        assert set(parsed) == {"entities", "relationships"}
        assert isinstance(parsed["entities"], list)
        assert isinstance(parsed["relationships"], list)
        assert parsed["entities"]
        assert parsed["relationships"]

        entity_names = {
            entity["name"] for entity in parsed.get("entities", []) if entity
        }
        for relationship in parsed.get("relationships", []):
            assert relationship["source"] in entity_names
            assert relationship["target"] in entity_names
        assert "<entity_name>" in entity_names


# ---------------------------------------------------------------------------
# 2. DEFAULT_ENTITY_TYPES is gone from constants
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_default_entity_types_removed_from_constants():
    """DEFAULT_ENTITY_TYPES must no longer exist in lightrag.constants."""
    import lightrag.constants as constants

    assert not hasattr(constants, "DEFAULT_ENTITY_TYPES"), (
        "DEFAULT_ENTITY_TYPES should have been removed from constants.py"
    )


# ---------------------------------------------------------------------------
# 3. ENTITY_TYPES env var raises SystemExit
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_entity_types_env_var_raises_system_exit(tmp_path):
    """LightRAG.__post_init__ must raise SystemExit when ENTITY_TYPES env var is set."""
    from lightrag import LightRAG

    with patch.dict(os.environ, {"ENTITY_TYPES": '["Person"]'}):
        with pytest.raises(SystemExit) as exc_info:
            LightRAG(
                working_dir=str(tmp_path),
                llm_model_func=AsyncMock(),
                embedding_func=None,
            )
    assert "ENTITY_TYPES" in str(exc_info.value)


# ---------------------------------------------------------------------------
# 4. Text mode: entity_types_guidance injected into prompt
# ---------------------------------------------------------------------------


@pytest.mark.offline
@pytest.mark.asyncio
async def test_text_mode_default_guidance_injected_into_prompt():
    """Default entity_types_guidance is passed to LLM system prompt in text mode."""
    from lightrag.operate import extract_entities
    from lightrag.prompt import PROMPTS

    global_config = _make_global_config(use_json=False)
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(
            chunks=_make_chunks(),
            global_config=global_config,
        )

    # The system prompt passed to the LLM must contain the default guidance
    assert llm_func.await_count >= 1
    call_kwargs = llm_func.call_args_list[0][1]
    system_prompt = call_kwargs.get("system_prompt", "")
    assert PROMPTS["default_entity_types_guidance"] in system_prompt
    assert "must start with `relation`, never `entity`" in system_prompt
    assert "After the last entity row, switch prefixes to `relation`" in system_prompt
    assert "Output at most 100 total rows" in system_prompt
    assert "Output at most 40 entity rows" in system_prompt


@pytest.mark.offline
@pytest.mark.asyncio
async def test_text_mode_custom_guidance_overrides_default():
    """Custom entity_types_guidance in addon_params overrides default."""
    from lightrag.operate import extract_entities

    custom_guidance = "- Widget: A test widget type"
    global_config = _make_global_config(
        addon_params={"entity_types_guidance": custom_guidance},
        use_json=False,
    )
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(
            chunks=_make_chunks(),
            global_config=global_config,
        )

    call_kwargs = llm_func.call_args_list[0][1]
    system_prompt = call_kwargs.get("system_prompt", "")
    assert custom_guidance in system_prompt


@pytest.mark.offline
def test_text_continue_prompt_requires_relation_prefix_for_corrections():
    from lightrag.prompt import PROMPTS

    prompt = PROMPTS["entity_continue_extraction_user_prompt"]
    assert (
        "Any corrected relationship row must be emitted with the literal `relation` prefix"
        in prompt
    )
    assert (
        "output at most {max_total_records} total rows and at most {max_entity_records} entity rows"
        in prompt
    )
    assert (
        "may reference entities that were already extracted correctly in the previous response"
        in prompt
    )
    assert (
        "whose source and target entities are both included in this response"
        not in prompt
    )


@pytest.mark.offline
def test_text_user_prompt_includes_quantity_limits():
    from lightrag.prompt import PROMPTS

    prompt = PROMPTS["entity_extraction_user_prompt"]
    assert (
        "output at most {max_total_records} total rows and at most {max_entity_records} entity rows"
        in prompt
    )
    assert (
        "If the row limit is reached, output `{completion_delimiter}` immediately"
        in prompt
    )


# ---------------------------------------------------------------------------
# 5. JSON mode: entity_types_guidance injected + entity_extraction kwarg set
# ---------------------------------------------------------------------------


@pytest.mark.offline
@pytest.mark.asyncio
async def test_rebuild_from_cached_fenced_json_uses_json_parser():
    """Cached JSON wrapped in markdown fences must not fall back to text parsing."""
    from lightrag import operate

    fenced_json = f"```json\n{_JSON_MODE_RESPONSE}\n```"

    with patch(
        "lightrag.operate._process_extraction_result",
        side_effect=AssertionError("text parser should not be used"),
    ):
        nodes, edges = await operate._rebuild_from_extraction_result(
            text_chunks_storage=_DummyTextChunksStorage(),
            extraction_result=fenced_json,
            chunk_id="chunk-001",
            timestamp=123,
        )

    assert set(nodes) == {"Alice", "Acme Corp"}
    assert ("Alice", "Acme Corp") in edges
    assert nodes["Alice"][0]["file_path"] == "test.md"


@pytest.mark.offline
@pytest.mark.asyncio
async def test_json_mode_default_guidance_injected_into_prompt():
    """Default entity_types_guidance is passed to LLM system prompt in JSON mode."""
    from lightrag.operate import extract_entities
    from lightrag.prompt import PROMPTS

    global_config = _make_global_config(use_json=True)
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _JSON_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(
            chunks=_make_chunks(),
            global_config=global_config,
        )

    assert llm_func.await_count >= 1
    call_kwargs = llm_func.call_args_list[0][1]
    system_prompt = call_kwargs.get("system_prompt", "")
    assert PROMPTS["default_entity_types_guidance"] in system_prompt
    assert "Output at most 100 total records" in system_prompt
    assert "Output at most 40 entity objects" in system_prompt


@pytest.mark.offline
@pytest.mark.asyncio
async def test_json_mode_entity_extraction_kwarg_passed():
    """JSON mode must pass response_format={'type':'json_object'} to the LLM function."""
    from lightrag.operate import extract_entities

    global_config = _make_global_config(use_json=True)
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _JSON_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(
            chunks=_make_chunks(),
            global_config=global_config,
        )

    assert llm_func.await_count >= 1
    call_kwargs = llm_func.call_args_list[0][1]
    assert call_kwargs.get("response_format") == {"type": "json_object"}
    assert call_kwargs.get("entity_extraction") is not True


@pytest.mark.offline
@pytest.mark.asyncio
async def test_json_mode_custom_guidance_overrides_default():
    """Custom entity_types_guidance in addon_params overrides default in JSON mode."""
    from lightrag.operate import extract_entities

    custom_guidance = "- Gadget: A test gadget type"
    global_config = _make_global_config(
        addon_params={"entity_types_guidance": custom_guidance},
        use_json=True,
    )
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _JSON_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(
            chunks=_make_chunks(),
            global_config=global_config,
        )

    call_kwargs = llm_func.call_args_list[0][1]
    system_prompt = call_kwargs.get("system_prompt", "")
    assert custom_guidance in system_prompt


@pytest.mark.offline
def test_json_user_prompt_includes_quantity_limits():
    from lightrag.prompt import PROMPTS

    prompt = PROMPTS["entity_extraction_json_user_prompt"]
    assert (
        "output at most {max_total_records} total records and at most {max_entity_records} entity objects"
        in prompt
    )
    assert (
        "Only output relationship objects whose `source` and `target` are both included"
        in prompt
    )


@pytest.mark.offline
def test_json_continue_prompt_includes_quantity_limits():
    from lightrag.prompt import PROMPTS

    prompt = PROMPTS["entity_continue_extraction_json_user_prompt"]
    assert (
        "output at most {max_total_records} total records and at most {max_entity_records} entity objects"
        in prompt
    )
    assert (
        "may reference entities already extracted correctly in the previous response"
        in prompt
    )


# ---------------------------------------------------------------------------
# 6. Text mode: entity_extraction kwarg NOT passed (only JSON mode uses it)
# ---------------------------------------------------------------------------


@pytest.mark.offline
@pytest.mark.asyncio
async def test_text_mode_no_entity_extraction_kwarg():
    """Text mode must NOT pass entity_extraction=True to the LLM function."""
    from lightrag.operate import extract_entities

    global_config = _make_global_config(use_json=False)
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(
            chunks=_make_chunks(),
            global_config=global_config,
        )

    call_kwargs = llm_func.call_args_list[0][1]
    assert not call_kwargs.get("entity_extraction", False)


@pytest.mark.offline
@pytest.mark.asyncio
async def test_text_mode_recovers_mis_prefixed_relationship_row():
    from lightrag.operate import extract_entities

    global_config = _make_global_config(use_json=False)
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_MISPREFIXED_RELATION_RESPONSE

    with patch("lightrag.operate.logger"):
        chunk_results = await extract_entities(
            chunks=_make_chunks(),
            global_config=global_config,
        )

    entities, relationships = chunk_results[0]
    assert len(entities) == 2
    assert len(relationships) == 1
    assert next(iter(relationships.keys())) == ("Alice", "Acme Corp")


@pytest.mark.offline
@pytest.mark.asyncio
async def test_text_mode_gleaned_relation_merges_cleanly_after_recovery():
    from lightrag.operate import extract_entities

    global_config = _make_global_config(use_json=False, max_gleaning=1)
    llm_func = global_config["llm_model_func"]
    llm_func.side_effect = _TEXT_MODE_GLEANED_RELATION_RESPONSES

    with patch("lightrag.operate.logger"):
        chunk_results = await extract_entities(
            chunks=_make_chunks(),
            global_config=global_config,
        )

    entities, relationships = chunk_results[0]
    assert len(entities) == 2
    assert len(relationships) == 1
    relation_data = next(iter(relationships.values()))[0]
    assert relation_data["src_id"] == "Alice"
    assert relation_data["tgt_id"] == "Acme Corp"


@pytest.mark.offline
@pytest.mark.asyncio
async def test_text_mode_gleaned_relation_can_reference_prior_entity():
    from lightrag.operate import extract_entities

    global_config = _make_global_config(use_json=False, max_gleaning=1)
    llm_func = global_config["llm_model_func"]
    llm_func.side_effect = _TEXT_MODE_CROSS_PASS_RELATION_RESPONSES

    with patch("lightrag.operate.logger"):
        chunk_results = await extract_entities(
            chunks=_make_chunks(),
            global_config=global_config,
        )

    entities, relationships = chunk_results[0]
    assert set(entities.keys()) == {"Alice", "Acme Corp"}
    assert len(relationships) == 1
    relation_data = next(iter(relationships.values()))[0]
    assert relation_data["src_id"] == "Alice"
    assert relation_data["tgt_id"] == "Acme Corp"


@pytest.mark.offline
def test_addon_params_default_includes_entity_type_prompt_file_env(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    prompt_dir = tmp_path / "entity_type"
    prompt_dir.mkdir()
    _write_prompt_profile(
        prompt_dir / "entity_type_prompt.sample.yml",
        text_examples=[_text_profile_example("Env Default Example")],
    )

    with patch.dict(
        os.environ,
        {
            "SUMMARY_LANGUAGE": "English",
            "ENTITY_TYPE_PROMPT_FILE": "entity_type_prompt.sample.yml",
        },
    ):
        with _patch_prompt_dir(prompt_dir):
            rag = LightRAG(
                working_dir=str(tmp_path / "rag-default-env"),
                llm_model_func=AsyncMock(),
                embedding_func=_dummy_embedding_func(),
                entity_extraction_use_json=False,
            )

    assert (
        rag.addon_params["entity_type_prompt_file"] == "entity_type_prompt.sample.yml"
    )


@pytest.mark.offline
@pytest.mark.asyncio
async def test_text_mode_prompt_file_injects_examples_and_guidance():
    _require_yaml()

    from lightrag.operate import extract_entities

    guidance = "- ExampleType: Injected guidance"
    example_label = "Custom Text Example"
    prompt_profile = {
        "entity_types_guidance": guidance,
        "entity_extraction_examples": [_text_profile_example(example_label)],
        "entity_extraction_json_examples": [],
    }

    global_config = _make_global_config(
        prompt_profile=prompt_profile,
        use_json=False,
    )
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(chunks=_make_chunks(), global_config=global_config)

    call_kwargs = llm_func.call_args_list[0][1]
    system_prompt = call_kwargs.get("system_prompt", "")
    assert guidance in system_prompt
    assert example_label in system_prompt


@pytest.mark.offline
@pytest.mark.asyncio
async def test_json_mode_prompt_file_injects_examples_and_guidance():
    _require_yaml()

    from lightrag.operate import extract_entities

    guidance = "- ExampleType: Injected JSON guidance"
    example_label = "Custom Json Example"
    prompt_profile = {
        "entity_types_guidance": guidance,
        "entity_extraction_examples": [],
        "entity_extraction_json_examples": [_json_profile_example(example_label)],
    }

    global_config = _make_global_config(
        prompt_profile=prompt_profile,
        use_json=True,
    )
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _JSON_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(chunks=_make_chunks(), global_config=global_config)

    call_kwargs = llm_func.call_args_list[0][1]
    system_prompt = call_kwargs.get("system_prompt", "")
    assert guidance in system_prompt
    assert example_label in system_prompt


@pytest.mark.offline
@pytest.mark.asyncio
async def test_prompt_file_guidance_falls_back_to_default_when_missing():
    _require_yaml()

    from lightrag.operate import extract_entities
    from lightrag.prompt import PROMPTS

    global_config = _make_global_config(
        prompt_profile={
            "entity_types_guidance": PROMPTS["default_entity_types_guidance"].rstrip(),
            "entity_extraction_examples": [
                _text_profile_example("Fallback Guidance Example")
            ],
            "entity_extraction_json_examples": [],
        },
        use_json=False,
    )
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(chunks=_make_chunks(), global_config=global_config)

    call_kwargs = llm_func.call_args_list[0][1]
    system_prompt = call_kwargs.get("system_prompt", "")
    assert PROMPTS["default_entity_types_guidance"] in system_prompt


@pytest.mark.offline
@pytest.mark.asyncio
async def test_cached_prompt_profile_supplies_merged_guidance():
    from lightrag.operate import extract_entities

    merged_guidance = "- ExampleType: Addon override"

    global_config = _make_global_config(
        prompt_profile={
            "entity_types_guidance": merged_guidance,
            "entity_extraction_examples": [_text_profile_example("Override Example")],
            "entity_extraction_json_examples": [],
        },
        use_json=False,
    )
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(chunks=_make_chunks(), global_config=global_config)

    call_kwargs = llm_func.call_args_list[0][1]
    system_prompt = call_kwargs.get("system_prompt", "")
    assert merged_guidance in system_prompt


@pytest.mark.offline
def test_text_mode_prompt_file_can_omit_json_examples(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    prompt_dir = tmp_path / "entity_type"
    prompt_dir.mkdir()
    _write_prompt_profile(
        prompt_dir / "text_only.yml",
        text_examples=[_text_profile_example("Text Only Example")],
    )

    with _patch_prompt_dir(prompt_dir):
        rag = LightRAG(
            working_dir=str(tmp_path / "rag-text"),
            llm_model_func=AsyncMock(),
            embedding_func=_dummy_embedding_func(),
            entity_extraction_use_json=False,
            addon_params={"entity_type_prompt_file": "text_only.yml"},
        )

    assert rag.addon_params["entity_type_prompt_file"] == "text_only.yml"


@pytest.mark.offline
def test_json_mode_prompt_file_can_omit_text_examples(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    prompt_dir = tmp_path / "entity_type"
    prompt_dir.mkdir()
    _write_prompt_profile(
        prompt_dir / "json_only.yml",
        json_examples=[_json_profile_example("Json Only Example")],
    )

    with _patch_prompt_dir(prompt_dir):
        rag = LightRAG(
            working_dir=str(tmp_path / "rag-json"),
            llm_model_func=AsyncMock(),
            embedding_func=_dummy_embedding_func(),
            entity_extraction_use_json=True,
            addon_params={"entity_type_prompt_file": "json_only.yml"},
        )

    assert rag.addon_params["entity_type_prompt_file"] == "json_only.yml"


@pytest.mark.offline
def test_text_mode_prompt_file_requires_text_examples(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    prompt_dir = tmp_path / "entity_type"
    prompt_dir.mkdir()
    _write_prompt_profile(
        prompt_dir / "missing_text_examples.yml",
        json_examples=[_json_profile_example("Wrong Mode Only")],
    )

    with _patch_prompt_dir(prompt_dir):
        with pytest.raises(ValueError) as exc_info:
            LightRAG(
                working_dir=str(tmp_path / "rag-missing-text"),
                llm_model_func=AsyncMock(),
                embedding_func=None,
                entity_extraction_use_json=False,
                addon_params={"entity_type_prompt_file": "missing_text_examples.yml"},
            )

    assert "entity_extraction_examples" in str(exc_info.value)


@pytest.mark.offline
def test_json_mode_prompt_file_requires_json_examples(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    prompt_dir = tmp_path / "entity_type"
    prompt_dir.mkdir()
    _write_prompt_profile(
        prompt_dir / "missing_json_examples.yml",
        text_examples=[_text_profile_example("Wrong Mode Only")],
    )

    with _patch_prompt_dir(prompt_dir):
        with pytest.raises(ValueError) as exc_info:
            LightRAG(
                working_dir=str(tmp_path / "rag-missing-json"),
                llm_model_func=AsyncMock(),
                embedding_func=None,
                entity_extraction_use_json=True,
                addon_params={"entity_type_prompt_file": "missing_json_examples.yml"},
            )

    assert "entity_extraction_json_examples" in str(exc_info.value)


@pytest.mark.offline
def test_prompt_file_rejects_directory_segments(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    with pytest.raises(ValueError) as exc_info:
        LightRAG(
            working_dir=str(tmp_path / "rag-bad-path"),
            llm_model_func=AsyncMock(),
            embedding_func=None,
            addon_params={"entity_type_prompt_file": "../outside.yml"},
        )

    assert "file name only" in str(exc_info.value)


@pytest.mark.offline
def test_prompt_file_rejects_absolute_paths(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    with pytest.raises(ValueError) as exc_info:
        LightRAG(
            working_dir=str(tmp_path / "rag-abs-path"),
            llm_model_func=AsyncMock(),
            embedding_func=None,
            addon_params={"entity_type_prompt_file": str(tmp_path / "abs.yml")},
        )

    assert "file name only" in str(exc_info.value)


@pytest.mark.offline
@pytest.mark.asyncio
async def test_extract_entities_uses_cached_prompt_profile_without_reloading():
    from lightrag.operate import extract_entities

    cached_profile = {
        "entity_types_guidance": "- ExampleType: Cached guidance",
        "entity_extraction_examples": [_text_profile_example("Cached Text Example")],
        "entity_extraction_json_examples": [],
    }
    global_config = _make_global_config(use_json=False, prompt_profile=cached_profile)
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    with patch(
        "lightrag.operate.resolve_entity_extraction_prompt_profile",
        side_effect=AssertionError("should not resolve profile when cache exists"),
    ):
        with patch("lightrag.operate.logger"):
            await extract_entities(chunks=_make_chunks(), global_config=global_config)
            await extract_entities(chunks=_make_chunks(), global_config=global_config)

    system_prompt = llm_func.call_args_list[0][1].get("system_prompt", "")
    assert "Cached Text Example" in system_prompt
    assert "Cached guidance" in system_prompt


@pytest.mark.offline
def test_sample_prompt_file_matches_builtin_prompt_data():
    _require_yaml()

    from lightrag.prompt import (
        get_default_entity_extraction_prompt_profile,
        load_entity_extraction_prompt_profile,
    )

    sample_file = (
        Path(__file__).resolve().parents[2]
        / "prompts"
        / "samples"
        / "entity_type_prompt.sample.yml"
    )

    loaded_profile = load_entity_extraction_prompt_profile(sample_file)
    assert loaded_profile == get_default_entity_extraction_prompt_profile()


@pytest.mark.offline
def test_prompt_dir_env_var_overrides_default(tmp_path, monkeypatch):
    _require_yaml()

    from lightrag.prompt import (
        get_entity_type_prompt_dir,
        resolve_entity_type_prompt_path,
    )

    monkeypatch.setenv("PROMPT_DIR", str(tmp_path))
    expected_dir = (tmp_path / "entity_type").resolve()
    assert get_entity_type_prompt_dir() == expected_dir
    resolved = resolve_entity_type_prompt_path("custom.yml")
    assert resolved == expected_dir / "custom.yml"


@pytest.mark.offline
def test_prompt_dir_defaults_to_cwd_relative(tmp_path, monkeypatch):
    _require_yaml()

    from lightrag.prompt import get_entity_type_prompt_dir

    monkeypatch.delenv("PROMPT_DIR", raising=False)
    monkeypatch.chdir(tmp_path)
    assert (
        get_entity_type_prompt_dir() == (tmp_path / "prompts" / "entity_type").resolve()
    )


@pytest.mark.offline
def test_prompt_file_rejects_unsupported_extension(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    with pytest.raises(ValueError, match="'.yml' or '.yaml'"):
        LightRAG(
            working_dir=str(tmp_path / "rag-bad-ext"),
            llm_model_func=AsyncMock(),
            embedding_func=None,
            addon_params={"entity_type_prompt_file": "profile.txt"},
        )


@pytest.mark.offline
def test_prompt_file_malformed_yaml_raises_valueerror(tmp_path):
    _require_yaml()

    from lightrag.prompt import load_entity_extraction_prompt_profile

    bad_file = tmp_path / "broken.yml"
    bad_file.write_text("entity_types_guidance: [unclosed", encoding="utf-8")

    with pytest.raises(ValueError, match="invalid YAML"):
        load_entity_extraction_prompt_profile(bad_file)


@pytest.mark.offline
def test_addon_guidance_overrides_file_profile(tmp_path):
    _require_yaml()

    from lightrag.prompt import resolve_entity_extraction_prompt_profile

    prompt_dir = tmp_path / "entity_type"
    prompt_dir.mkdir()
    _write_prompt_profile(
        prompt_dir / "profile.yml",
        guidance="- FileType: from file",
        text_examples=[_text_profile_example("Merged Example")],
    )

    with _patch_prompt_dir(prompt_dir):
        profile = resolve_entity_extraction_prompt_profile(
            addon_params={
                "entity_type_prompt_file": "profile.yml",
                "entity_types_guidance": "- AddonType: from addon_params",
            },
            use_json=False,
        )

    assert profile["entity_types_guidance"] == "- AddonType: from addon_params"
    # File-provided examples must still be honored.
    assert any(
        "Merged Example" in example for example in profile["entity_extraction_examples"]
    )


@pytest.mark.offline
def test_explicit_addon_params_still_picks_up_env_defaults(tmp_path, monkeypatch):
    """Passing addon_params explicitly must not drop env-based defaults."""
    _require_yaml()

    from lightrag import LightRAG

    prompt_dir = tmp_path / "entity_type"
    prompt_dir.mkdir()
    _write_prompt_profile(
        prompt_dir / "from_env.yml",
        text_examples=[_text_profile_example("Env Example")],
    )

    monkeypatch.setenv("ENTITY_TYPE_PROMPT_FILE", "from_env.yml")

    with _patch_prompt_dir(prompt_dir):
        rag = LightRAG(
            working_dir=str(tmp_path / "rag-env-default"),
            llm_model_func=AsyncMock(),
            embedding_func=_dummy_embedding_func(),
            entity_extraction_use_json=False,
            addon_params={"language": "English"},
        )

    assert rag.addon_params["entity_type_prompt_file"] == "from_env.yml"


@pytest.mark.offline
def test_runtime_addon_params_item_update_refreshes_cached_values(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    prompt_dir = tmp_path / "entity_type"
    prompt_dir.mkdir()
    _write_prompt_profile(
        prompt_dir / "initial.yml",
        text_examples=[_text_profile_example("Initial Example")],
    )
    _write_prompt_profile(
        prompt_dir / "updated.yml",
        guidance="- UpdatedType: runtime update",
        text_examples=[_text_profile_example("Updated Example")],
    )

    with _patch_prompt_dir(prompt_dir):
        rag = LightRAG(
            working_dir=str(tmp_path / "rag-runtime-update"),
            llm_model_func=AsyncMock(),
            embedding_func=_dummy_embedding_func(),
            entity_extraction_use_json=False,
            addon_params={
                "entity_type_prompt_file": "initial.yml",
                "language": "English",
            },
        )
        rag.addon_params["entity_type_prompt_file"] = "updated.yml"
        rag.addon_params["language"] = "French"
        global_config = rag._build_global_config()

    assert global_config["addon_params"]["language"] == "French"
    assert global_config["_resolved_summary_language"] == "French"
    assert (
        global_config["_entity_extraction_prompt_profile"]["entity_types_guidance"]
        == "- UpdatedType: runtime update"
    )
    assert any(
        "Updated Example" in example
        for example in global_config["_entity_extraction_prompt_profile"][
            "entity_extraction_examples"
        ]
    )


@pytest.mark.offline
def test_runtime_addon_params_replacement_refreshes_cached_values(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    rag = LightRAG(
        working_dir=str(tmp_path / "rag-runtime-replace"),
        llm_model_func=AsyncMock(),
        embedding_func=_dummy_embedding_func(),
        entity_extraction_use_json=False,
        addon_params={"language": "English"},
    )

    rag.addon_params = {
        "language": "German",
        "entity_types_guidance": "- ReplacementType: runtime replace",
    }
    global_config = rag._build_global_config()

    assert global_config["addon_params"]["language"] == "German"
    assert global_config["_resolved_summary_language"] == "German"
    assert (
        global_config["_entity_extraction_prompt_profile"]["entity_types_guidance"]
        == "- ReplacementType: runtime replace"
    )


@pytest.mark.offline
def test_runtime_mode_flip_invalidates_cached_prompt_profile(tmp_path):
    _require_yaml()

    from lightrag import LightRAG

    prompt_dir = tmp_path / "entity_type"
    prompt_dir.mkdir()
    _write_prompt_profile(
        prompt_dir / "text_only.yml",
        text_examples=[_text_profile_example("Text Only Example")],
    )

    with _patch_prompt_dir(prompt_dir):
        rag = LightRAG(
            working_dir=str(tmp_path / "rag-mode-flip"),
            llm_model_func=AsyncMock(),
            embedding_func=_dummy_embedding_func(),
            entity_extraction_use_json=False,
            addon_params={"entity_type_prompt_file": "text_only.yml"},
        )

        rag._build_global_config()
        rag.entity_extraction_use_json = True

        with pytest.raises(ValueError) as exc_info:
            rag._build_global_config()

    assert "entity_extraction_json_examples" in str(exc_info.value)


# ---------------------------------------------------------------------------
# Section Context (heading breadcrumb) injection into extraction user prompts
# ---------------------------------------------------------------------------

_SECTION_MARKER = "---Section Context---"


def _render_text_user_prompt(heading_context_block: str) -> str:
    from lightrag.prompt import PROMPTS

    return PROMPTS["entity_extraction_user_prompt"].format(
        max_total_records=100,
        max_entity_records=40,
        completion_delimiter="<|COMPLETE|>",
        language="English",
        input_text="Alice founded Acme Corp.",
        heading_context_block=heading_context_block,
    )


def _render_json_user_prompt(heading_context_block: str) -> str:
    from lightrag.prompt import PROMPTS

    return PROMPTS["entity_extraction_json_user_prompt"].format(
        max_total_records=100,
        max_entity_records=40,
        language="English",
        entity_types_guidance="- Person: humans",
        input_text="Alice founded Acme Corp.",
        heading_context_block=heading_context_block,
    )


def _section_block(heading_path: str) -> str:
    from lightrag.prompt import PROMPTS

    return PROMPTS["entity_extraction_section_context"].format(
        heading_path=heading_path
    )


@pytest.mark.offline
def test_user_prompts_keep_single_real_input_text_section():
    """Only the rendered task prompt should carry the real input section marker."""

    text_markers = [
        line
        for line in _render_text_user_prompt("").splitlines()
        if line == "---Input Text---"
    ]
    json_markers = [
        line
        for line in _render_json_user_prompt("").splitlines()
        if line == "---Input Text---"
    ]

    assert len(text_markers) == 1
    assert len(json_markers) == 1


@pytest.mark.offline
def test_format_heading_context_full_path_includes_current_heading():
    """The breadcrumb appends the chunk's own heading after the parent chain."""
    from lightrag.chunk_schema import format_heading_context

    chunk = {
        "content": "...",
        "heading": {
            "level": 2,
            "heading": "Data Collection",
            "parent_headings": ["Methods"],
        },
    }
    assert format_heading_context(chunk) == "Methods → Data Collection"


@pytest.mark.offline
def test_format_heading_context_empty_when_no_heading():
    """A chunk without heading info yields an empty breadcrumb (block omitted)."""
    from lightrag.chunk_schema import format_heading_context

    chunk = {
        "content": "...",
        "tokens": 1,
        "full_doc_id": "d",
        "chunk_order_index": 0,
    }
    assert format_heading_context(chunk) == ""


@pytest.mark.offline
def test_text_user_prompt_section_context_hidden_and_byte_identical_when_no_heading():
    """No heading -> the whole `---Section Context---` block disappears and the
    rendered text user prompt is byte-identical to the placeholder-free form."""
    from lightrag.prompt import PROMPTS

    rendered = _render_text_user_prompt("")
    assert _SECTION_MARKER not in rendered

    # The placeholder is the ONLY change to this template, so rendering it empty
    # must equal a version with the placeholder physically removed (i.e. the
    # pre-change template). This is the hard no-noise regression guard.
    baseline_template = PROMPTS["entity_extraction_user_prompt"].replace(
        "{heading_context_block}", ""
    )
    baseline = baseline_template.format(
        max_total_records=100,
        max_entity_records=40,
        completion_delimiter="<|COMPLETE|>",
        language="English",
        input_text="Alice founded Acme Corp.",
    )
    assert rendered == baseline


@pytest.mark.offline
def test_json_user_prompt_section_context_hidden_and_byte_identical_when_no_heading():
    from lightrag.prompt import PROMPTS

    rendered = _render_json_user_prompt("")
    assert _SECTION_MARKER not in rendered

    baseline_template = PROMPTS["entity_extraction_json_user_prompt"].replace(
        "{heading_context_block}", ""
    )
    baseline = baseline_template.format(
        max_total_records=100,
        max_entity_records=40,
        language="English",
        entity_types_guidance="- Person: humans",
        input_text="Alice founded Acme Corp.",
    )
    assert rendered == baseline


@pytest.mark.offline
def test_text_user_prompt_includes_section_context_when_heading_present():
    rendered = _render_text_user_prompt(_section_block("Methods → Data Collection"))
    assert _SECTION_MARKER in rendered
    assert "Methods → Data Collection" in rendered
    # Block sits immediately above the input text section.
    assert "Methods → Data Collection\n\n---Input Text---" in rendered


@pytest.mark.offline
def test_json_user_prompt_includes_section_context_when_heading_present():
    rendered = _render_json_user_prompt(_section_block("Methods → Data Collection"))
    assert _SECTION_MARKER in rendered
    assert "Methods → Data Collection" in rendered
    assert "Methods → Data Collection\n\n---Input Text---" in rendered


@pytest.mark.offline
def test_section_context_breadcrumb_is_not_at_line_start():
    """A heading that looks like a prompt marker must be rendered inline (as
    data), never at the start of a line where it could forge a new section."""
    block = _section_block("---Output---")
    # The breadcrumb follows a label on the same line, so the marker text never
    # begins a line of its own.
    assert "\n---Output---" not in block
    assert "---Output---" in block  # still present, just inert/inline


@pytest.mark.offline
def test_extraction_system_prompts_reference_section_context():
    """Both system prompts carry the static conditional instruction."""
    from lightrag.prompt import PROMPTS

    for key in (
        "entity_extraction_system_prompt",
        "entity_extraction_json_system_prompt",
    ):
        assert _SECTION_MARKER in PROMPTS[key]
        assert "only as background" in PROMPTS[key]


@pytest.mark.offline
@pytest.mark.asyncio
async def test_extract_entities_injects_section_context_for_chunk_with_heading():
    """End-to-end: a chunk carrying a heading produces a user prompt containing
    its full section breadcrumb; a heading-free chunk does not."""
    from lightrag.operate import extract_entities

    global_config = _make_global_config(use_json=False)
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    chunks = {
        "chunk-001": {
            "tokens": 10,
            "content": "Alice founded Acme Corp.",
            "full_doc_id": "doc-001",
            "chunk_order_index": 0,
            "heading": {
                "level": 2,
                "heading": "Data Collection",
                "parent_headings": ["Methods"],
            },
        }
    }

    with patch("lightrag.operate.logger"):
        await extract_entities(chunks=chunks, global_config=global_config)

    assert llm_func.await_count >= 1
    user_prompt = llm_func.call_args_list[0][0][0]
    assert _SECTION_MARKER in user_prompt
    assert "Methods → Data Collection\n\n---Input Text---" in user_prompt


@pytest.mark.offline
@pytest.mark.asyncio
async def test_extract_entities_omits_section_context_for_chunk_without_heading():
    from lightrag.operate import extract_entities

    global_config = _make_global_config(use_json=False)
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    with patch("lightrag.operate.logger"):
        await extract_entities(chunks=_make_chunks(), global_config=global_config)

    assert llm_func.await_count >= 1
    user_prompt = llm_func.call_args_list[0][0][0]
    assert _SECTION_MARKER not in user_prompt


# ---------------------------------------------------------------------------
# Section Context length bounding: per-level char cap + overall token budget
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_format_heading_context_caps_long_level():
    """A single runaway heading level is truncated to the per-level char cap."""
    from lightrag.chunk_schema import (
        DEFAULT_HEADING_LEVEL_MAX_CHARS,
        format_heading_context,
    )

    long_title = "A" * (DEFAULT_HEADING_LEVEL_MAX_CHARS + 50)
    chunk = {"heading": {"level": 1, "heading": long_title, "parent_headings": []}}

    out = format_heading_context(chunk)
    assert out.endswith("…")
    assert len(out) == DEFAULT_HEADING_LEVEL_MAX_CHARS


@pytest.mark.offline
def test_format_heading_context_per_level_cap_can_be_disabled():
    from lightrag.chunk_schema import format_heading_context

    long_title = "B" * 300
    chunk = {"heading": {"level": 1, "heading": long_title, "parent_headings": []}}

    assert format_heading_context(chunk, max_heading_len=0) == long_title


# ---------------------------------------------------------------------------
# Query-stage format_parent_headings: same per-level cap + cleaning as extraction
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_format_parent_headings_caps_long_level():
    """A runaway parent heading is truncated to the per-level char cap, matching
    the extraction-stage format_heading_context."""
    from lightrag.chunk_schema import (
        DEFAULT_HEADING_LEVEL_MAX_CHARS,
        format_parent_headings,
    )

    long_title = "A" * (DEFAULT_HEADING_LEVEL_MAX_CHARS + 50)
    chunk = {
        "heading": {"level": 2, "heading": "Leaf", "parent_headings": [long_title]}
    }

    out = format_parent_headings(chunk)
    assert out.endswith("…")
    assert len(out) == DEFAULT_HEADING_LEVEL_MAX_CHARS  # only the parent, capped


@pytest.mark.offline
def test_format_parent_headings_per_level_cap_can_be_disabled():
    from lightrag.chunk_schema import format_parent_headings

    long_title = "B" * 300
    chunk = {
        "heading": {"level": 2, "heading": "Leaf", "parent_headings": [long_title]}
    }

    assert format_parent_headings(chunk, max_heading_len=0) == long_title


@pytest.mark.offline
def test_format_parent_headings_cleaning_matches_extraction():
    """Parent headings get the same cleaning as extraction: → folded to a space,
    Cc/Cf control chars stripped (shared _clean_heading_text)."""
    from lightrag.chunk_schema import format_parent_headings

    # chr(0) is a Cc control; chr(0x200B) is ZWSP (Cf) — both stripped. Built
    # via chr() so the source carries no literal invisible characters.
    second_level = "x" + chr(0) + "y" + chr(0x200B) + "z"
    chunk = {
        "heading": {
            "level": 2,
            "heading": "Leaf",
            "parent_headings": ["A→B", second_level],
        }
    }
    # "A→B" -> "A B"; control + format chars removed from the second level.
    assert format_parent_headings(chunk) == "A B → xyz"


@pytest.mark.offline
def test_format_parent_headings_basic_behavior_preserved():
    """Existing behavior is unchanged: empty when no heading, normal multi-level
    path joined with the breadcrumb separator."""
    from lightrag.chunk_schema import format_parent_headings

    assert format_parent_headings({"content": "...", "chunk_order_index": 0}) == ""

    chunk = {
        "heading": {"level": 2, "heading": "Leaf", "parent_headings": ["h1", "h2"]}
    }
    assert format_parent_headings(chunk) == "h1 → h2"  # leaf NOT appended


class _FakeChunksDB:
    """Minimal text_chunks_db for _attach_content_headings: get_by_ids + config."""

    def __init__(self, data_by_id: dict, tokenizer):
        self._data = data_by_id
        self.global_config = {"tokenizer": tokenizer}

    async def get_by_ids(self, ids):
        return [self._data.get(i) for i in ids]


@pytest.mark.offline
@pytest.mark.asyncio
async def test_attach_content_headings_token_budgets_deep_path():
    """A deep heading chain (per-level cap bounds length, not count) is collapsed
    to fit DEFAULT_MAX_SECTION_CONTEXT_TOKENS, mirroring the extraction stage."""
    from lightrag.chunk_schema import HEADING_BREADCRUMB_SEP
    from lightrag.constants import DEFAULT_MAX_SECTION_CONTEXT_TOKENS
    from lightrag.operate import _attach_content_headings

    tok = Tokenizer("dummy", DummyTokenizer())  # 1 char == 1 token
    deep = [f"Level{i:02d}" for i in range(100)]  # well over the token budget
    db = _FakeChunksDB(
        {"c1": {"heading": {"level": 99, "heading": "Leaf", "parent_headings": deep}}},
        tok,
    )
    chunks = [{"chunk_id": "c1"}]

    await _attach_content_headings(chunks, db)

    out = chunks[0]["content_headings"]
    assert len(tok.encode(out)) <= DEFAULT_MAX_SECTION_CONTEXT_TOKENS
    # Collapsed to first → … → leaf, so a middle level is gone.
    assert f"{HEADING_BREADCRUMB_SEP}…{HEADING_BREADCRUMB_SEP}" in out
    assert "Level50" not in out


@pytest.mark.offline
@pytest.mark.asyncio
async def test_attach_content_headings_keeps_short_path_intact():
    """A within-budget path is attached unchanged (no token collapsing)."""
    from lightrag.operate import _attach_content_headings

    tok = Tokenizer("dummy", DummyTokenizer())
    db = _FakeChunksDB(
        {
            "c1": {
                "heading": {
                    "level": 2,
                    "heading": "Leaf",
                    "parent_headings": ["h1", "h2"],
                }
            }
        },
        tok,
    )
    chunks = [{"chunk_id": "c1"}]

    await _attach_content_headings(chunks, db)

    assert chunks[0]["content_headings"] == "h1 → h2"


@pytest.mark.offline
def test_truncate_section_context_noop_within_budget():
    from lightrag.operate import _truncate_section_context

    tok = Tokenizer("dummy", DummyTokenizer())
    path = "Methods → Data Collection"
    assert _truncate_section_context(path, tok, 256) == path


@pytest.mark.offline
def test_truncate_section_context_keeps_first_and_last_when_over_budget():
    """Over budget -> keep first (top-level) + last (leaf) section, elide middle."""
    from lightrag.chunk_schema import HEADING_BREADCRUMB_SEP
    from lightrag.operate import _truncate_section_context

    tok = Tokenizer("dummy", DummyTokenizer())  # 1 char == 1 token
    levels = [f"Level{i:02d}" for i in range(100)]
    path = HEADING_BREADCRUMB_SEP.join(levels)
    # Budget large enough for the collapsed two-level form (~21 tokens) so the
    # hard-cap backstop does not also fire here.
    budget = 40

    out = _truncate_section_context(path, tok, budget)
    expected = (
        f"{levels[0]}{HEADING_BREADCRUMB_SEP}…{HEADING_BREADCRUMB_SEP}{levels[-1]}"
    )
    assert out == expected
    assert "Level50" not in out  # middle levels are gone
    assert len(tok.encode(out)) <= budget


@pytest.mark.offline
def test_truncate_section_context_hard_caps_dense_short_path():
    """A 1-/2-level path that is itself over budget must still be capped
    (not bypassed) — guards token-dense / byte-level tokenizers."""
    from lightrag.chunk_schema import HEADING_BREADCRUMB_SEP
    from lightrag.operate import _truncate_section_context

    tok = Tokenizer("dummy", DummyTokenizer())
    path = HEADING_BREADCRUMB_SEP.join(["A" * 50, "B" * 50])  # 103 chars/tokens
    budget = 10

    out = _truncate_section_context(path, tok, budget)
    assert out != path
    assert out.endswith("…")
    assert len(tok.encode(out)) <= budget


@pytest.mark.offline
def test_truncate_section_context_accounts_for_multitoken_ellipsis():
    """The hard cap must reserve the tokenizer's actual ellipsis cost."""
    from lightrag.operate import _truncate_section_context

    class TwoTokenEllipsisTokenizer(TokenizerInterface):
        def encode(self, content: str):
            tokens = []
            for ch in content:
                if ch == "…":
                    tokens.extend([0x110000, 0x110001])
                else:
                    tokens.append(ord(ch))
            return tokens

        def decode(self, tokens):
            return "".join(chr(token) for token in tokens if token <= 0x10FFFF)

    tok = Tokenizer("two-token-ellipsis", TwoTokenEllipsisTokenizer())
    budget = 10

    out = _truncate_section_context("A" * 20, tok, budget)
    assert out == "A" * 8 + "…"
    assert len(tok.encode(out)) <= budget


@pytest.mark.offline
def test_truncate_section_context_hard_caps_collapsed_form_when_still_over():
    """Even the collapsed first→…→leaf form is capped if it still exceeds."""
    from lightrag.chunk_schema import HEADING_BREADCRUMB_SEP
    from lightrag.operate import _truncate_section_context

    tok = Tokenizer("dummy", DummyTokenizer())
    levels = [f"Level{i:02d}" for i in range(100)]
    path = HEADING_BREADCRUMB_SEP.join(levels)
    budget = 8  # smaller than the ~21-token collapsed form

    out = _truncate_section_context(path, tok, budget)
    assert out.endswith("…")
    assert len(tok.encode(out)) <= budget


@pytest.mark.offline
def test_heading_level_cap_below_one_third_of_token_budget():
    """Invariant guard: collapsed first+leaf must fit the token budget."""
    from lightrag.constants import (
        DEFAULT_HEADING_LEVEL_MAX_CHARS,
        DEFAULT_MAX_SECTION_CONTEXT_TOKENS,
    )

    assert DEFAULT_HEADING_LEVEL_MAX_CHARS * 3 < DEFAULT_MAX_SECTION_CONTEXT_TOKENS


@pytest.mark.offline
def test_truncate_section_context_disabled_or_no_tokenizer():
    from lightrag.operate import _truncate_section_context

    tok = Tokenizer("dummy", DummyTokenizer())
    path = "X" * 1000
    assert _truncate_section_context(path, tok, 0) == path
    assert _truncate_section_context(path, None, 256) == path


@pytest.mark.offline
@pytest.mark.asyncio
async def test_extract_entities_bounds_pathological_heading_in_prompt():
    """A chunk with an absurdly long heading must not inject it verbatim."""
    from lightrag.chunk_schema import DEFAULT_HEADING_LEVEL_MAX_CHARS
    from lightrag.operate import extract_entities

    global_config = _make_global_config(use_json=False)
    llm_func = global_config["llm_model_func"]
    llm_func.return_value = _TEXT_MODE_RESPONSE

    long_title = "Z" * 500
    chunks = {
        "chunk-001": {
            "tokens": 10,
            "content": "Alice founded Acme Corp.",
            "full_doc_id": "doc-001",
            "chunk_order_index": 0,
            "heading": {
                "level": 1,
                "heading": long_title,
                "parent_headings": [],
            },
        }
    }

    with patch("lightrag.operate.logger"):
        await extract_entities(chunks=chunks, global_config=global_config)

    user_prompt = llm_func.call_args_list[0][0][0]
    assert _SECTION_MARKER in user_prompt
    assert long_title not in user_prompt  # full title never reaches the prompt
    assert "Z" * DEFAULT_HEADING_LEVEL_MAX_CHARS not in user_prompt


# ---------------------------------------------------------------------------
# Heading text symbol cleaning: → -> space, strip Cc/Cf, preserve everything else
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_clean_heading_text_converts_arrow_to_space():
    """The breadcrumb separator char must never survive inside one heading."""
    from lightrag.chunk_schema import _clean_heading_text

    assert _clean_heading_text("A→B") == "A B"
    assert _clean_heading_text("A  →  B") == "A B"


@pytest.mark.offline
def test_clean_heading_text_strips_control_and_format_chars():
    """Cc (NUL, BEL, file/unit separators) and Cf (zero-width marks) are removed."""
    from lightrag.chunk_schema import _clean_heading_text

    # \x00 (Cc), ​ ZWSP (Cf),  BOM (Cf) all vanish.
    assert _clean_heading_text("a\x00b​c") == "abc"
    assert _clean_heading_text("x\x07y") == "xy"
    # \x1c-\x1f are Cc but NOT matched by \s — must be stripped, not kept.
    assert _clean_heading_text("p\x1c\x1fq") == "pq"


@pytest.mark.offline
def test_clean_heading_text_preserves_normal_characters():
    """CJK / Latin / digits / punctuation are left untouched; only → is folded."""
    from lightrag.chunk_schema import _clean_heading_text

    assert _clean_heading_text("方法 → 数据采集 (2024)!") == "方法 数据采集 (2024)!"
    # Adjacent CJK never gets a space inserted between characters.
    assert _clean_heading_text("数据采集") == "数据采集"


@pytest.mark.offline
def test_clean_heading_text_whitespace_collapse_is_last():
    """Newline/tab still fold to a single space (kept through the strip pass)."""
    from lightrag.chunk_schema import _clean_heading_text

    assert _clean_heading_text("a\nb\tc") == "a b c"
    # A control char removed between two words must not leave a double space.
    assert _clean_heading_text("a \x00 b") == "a b"


@pytest.mark.offline
def test_format_heading_context_arrow_in_heading_does_not_forge_level():
    """A heading containing → is cleaned, so the breadcrumb split stays accurate."""
    from lightrag.chunk_schema import (
        HEADING_BREADCRUMB_SEP,
        format_heading_context,
    )

    chunk = {
        "heading": {"level": 2, "heading": "C", "parent_headings": ["A→B"]},
    }
    out = format_heading_context(chunk)
    assert out == "A B → C"
    # The breadcrumb still splits into exactly the two real levels.
    assert out.split(HEADING_BREADCRUMB_SEP) == ["A B", "C"]