"""Prompt construction for LLM-NER extraction.

:class:`PromptBuilder` assembles the complete prompt from:

* A *role* description (who the LLM is).
* A *task* description (what the LLM must do).
* Per-field extraction rules collected by :class:`~llmner.factories.SchemaRegistry`.
* A JSON schema skeleton generated by :meth:`~llmner.base_model.NERBaseModel.prompt_schema`.
* The *input text* to analyse.

The default template is English-language and model-agnostic.  Pass a custom
``prompt_template`` string to :class:`PromptBuilder` to adapt it for a
specific language or model family.
"""

from __future__ import annotations

import textwrap

from llmner.base_model import NERBaseModel

__all__ = ["PromptBuilder", "DEFAULT_PROMPT_TEMPLATE"]

DEFAULT_PROMPT_TEMPLATE = """\
ENTITY EXTRACTION INSTRUCTION

ROLE: {system_role}

TASK: {system_task}

FIELD-SPECIFIC RULES:
{rules_text}

OUTPUT SCHEMA (replicate these exact keys, every terminal value must be \
{{ "value": ..., "evidence": ... }}):
{schema_json}

TEXT TO PROCESS:
"{input_text}"
"""


class PromptBuilder:
    """Assembles the complete LLM prompt for a given input text.

    Args:
        schema_class:     The :class:`~llmner.base_model.NERBaseModel`
                          subclass whose :meth:`~llmner.base_model.NERBaseModel.prompt_schema`
                          is used to build the JSON skeleton.
        system_role:      A brief description of the LLM's role
                          (e.g. ``"You are a forensic data analyst..."``).
        system_task:      A description of the extraction task and any
                          global constraints.
        rules_registry:   List of per-field rules produced by
                          :class:`~llmner.factories.SchemaRegistry`.
        prompt_template:  Optional override for the full prompt template.
                          Must contain ``{system_role}``, ``{system_task}``,
                          ``{rules_text}``, ``{schema_json}``, and
                          ``{input_text}`` placeholders.
    """

    def __init__(
        self,
        schema_class: type[NERBaseModel],
        system_role: str,
        system_task: str,
        rules_registry: list[str],
        prompt_template: str = DEFAULT_PROMPT_TEMPLATE,
    ) -> None:
        self._schema_class = schema_class
        self._system_role = system_role
        self._system_task = system_task
        self._rules_registry = rules_registry
        self._template = prompt_template

    def build(self, input_text: str) -> str:
        """Return the complete prompt string for *input_text*.

        Args:
            input_text: The source text to extract entities from.

        Returns:
            A fully populated prompt string ready to send to the LLM.
        """
        rules_text = "\n".join(f"- {r}" for r in self._rules_registry)
        schema_json = self._schema_class.prompt_schema()

        raw = self._template.format(
            system_role=self._system_role,
            system_task=self._system_task,
            rules_text=rules_text,
            schema_json=schema_json,
            input_text=input_text,
        )
        return textwrap.dedent(raw).strip()
