"""Unit tests for :mod:`llmner.prompt`."""

from __future__ import annotations

from llmner import NERBaseModel, SchemaRegistry
from llmner.prompt import DEFAULT_PROMPT_TEMPLATE, PromptBuilder

# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------

_reg = SchemaRegistry()
_NameType = _reg.generic("name", "Full name.")


class _TestSchema(NERBaseModel):
    name: _NameType | None = None  # type: ignore[valid-type]


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestPromptBuilder:
    def test_build_contains_all_sections(self) -> None:
        pb = PromptBuilder(
            schema_class=_TestSchema,
            system_role="You are a tester.",
            system_task="Extract names.",
            rules_registry=_reg.rules,
        )
        prompt = pb.build("John went to the store.")

        assert "You are a tester." in prompt
        assert "Extract names." in prompt
        assert "NAME:" in prompt
        assert "John went to the store." in prompt
        assert '"name"' in prompt  # schema JSON key

    def test_custom_template(self) -> None:
        custom = "ROLE={system_role} TASK={system_task} RULES={rules_text} SCHEMA={schema_json} TEXT={input_text}"
        pb = PromptBuilder(
            schema_class=_TestSchema,
            system_role="R",
            system_task="T",
            rules_registry=["rule1"],
            prompt_template=custom,
        )
        prompt = pb.build("hello")

        assert "ROLE=R" in prompt
        assert "TASK=T" in prompt
        assert "rule1" in prompt
        assert "hello" in prompt

    def test_input_text_with_braces_does_not_crash(self) -> None:
        """Curly braces in input text must not be interpreted as format keys."""
        pb = PromptBuilder(
            schema_class=_TestSchema,
            system_role="R",
            system_task="T",
            rules_registry=_reg.rules,
        )
        # This should not raise KeyError / ValueError
        prompt = pb.build('He said {"key": "value"} loudly.')
        assert '{"key": "value"}' in prompt

    def test_default_template_has_required_placeholders(self) -> None:
        for placeholder in [
            "{system_role}",
            "{system_task}",
            "{rules_text}",
            "{schema_json}",
            "{input_text}",
        ]:
            assert placeholder in DEFAULT_PROMPT_TEMPLATE
