"""Main extractor orchestrating prompt building, LLM calls, and parsing.

:class:`NERExtractor` is the primary entry-point for users of the library.
It combines a :class:`~llmner.base_model.NERBaseModel` schema with a
:class:`~llmner.llm_client.BaseLLMClient` and a
:class:`~llmner.prompt.PromptBuilder` to extract structured entities from
free text.

Quick start::

    from llmner import NERExtractor, NERBaseModel, SchemaRegistry

    registry = SchemaRegistry()

    class PersonSchema(NERBaseModel):
        name: registry.generic("name", "Full name") = None
        age: registry.int_range("age", "Age in years") = None

    extractor = NERExtractor(
        schema_class=PersonSchema,
        system_role="You are an expert information extractor.",
        system_task=(
            "Extract the requested fields from the text. "
            "Return null for any field not mentioned."
        ),
        rules_registry=registry.rules,
    )

    result = extractor.extract_one("John Smith is 34 years old.")
    print(result.name.value)  # "John Smith"
    print(result.age.value)   # "34"
"""

from __future__ import annotations

import logging

from llmner.base_model import NERBaseModel
from llmner.llm_client import BaseLLMClient, OllamaClient
from llmner.prompt import DEFAULT_PROMPT_TEMPLATE, PromptBuilder

__all__ = ["NERExtractor"]

logger = logging.getLogger(__name__)


class NERExtractor:
    """Orchestrates entity extraction from a single text block.

    Args:
        schema_class:     A :class:`~llmner.base_model.NERBaseModel` subclass
                          defining the fields to extract.
        system_role:      Brief description of the LLM's persona / expertise.
        system_task:      Full description of the extraction task and
                          constraints sent to the LLM.
        rules_registry:   Per-field extraction rules collected by
                          :class:`~llmner.factories.SchemaRegistry`.
        prompt_template:  Optional custom prompt template string.
        llm_client:       An instance of
                          :class:`~llmner.llm_client.BaseLLMClient`.
                          Defaults to :class:`~llmner.llm_client.OllamaClient`
                          with ``llm_model`` / ``llm_base_url`` /
                          ``llm_timeout``.
        llm_base_url:     Ollama server URL (ignored when *llm_client* is
                          provided).
        llm_model:        Ollama model tag (ignored when *llm_client* is
                          provided).
        llm_timeout:      Request timeout in seconds (ignored when
                          *llm_client* is provided).
        max_retries:      How many additional LLM calls to perform when
                          :meth:`~llmner.base_model.NERBaseModel.has_missing_fields`
                          returns ``True``.  Results are merged via
                          :meth:`~llmner.base_model.NERBaseModel.merge`.
                          Default is ``1``.
    """

    def __init__(
        self,
        schema_class: type[NERBaseModel],
        system_role: str,
        system_task: str,
        rules_registry: list[str],
        prompt_template: str = DEFAULT_PROMPT_TEMPLATE,
        *,
        llm_client: BaseLLMClient | None = None,
        llm_base_url: str = "http://localhost:11434",
        llm_model: str = "qwen2.5:7b-instruct",
        llm_timeout: int = 120,
        max_retries: int = 1,
    ) -> None:
        self.schema_class = schema_class
        self.max_retries = max_retries

        self._prompt_builder = PromptBuilder(
            schema_class=schema_class,
            system_role=system_role,
            system_task=system_task,
            rules_registry=rules_registry,
            prompt_template=prompt_template,
        )
        self._llm_client: BaseLLMClient = llm_client or OllamaClient(
            base_url=llm_base_url,
            model=llm_model,
            timeout=llm_timeout,
        )

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def extract_one(
        self,
        input_text: str,
        *,
        retry_on_null: bool = True,
    ) -> NERBaseModel | None:
        """Extract entities from a single text block.

        Args:
            input_text:    Source text to analyse.
            retry_on_null: When ``True`` (default), retry up to
                           ``self.max_retries`` times if any extracted value
                           is ``None``.  Each retry result is merged with the
                           previous one so that fields already extracted are
                           not lost.

        Returns:
            Populated :class:`~llmner.base_model.NERBaseModel` instance, or
            ``None`` if the LLM returned an unparseable response.
        """
        extraction = self._call_llm(input_text)
        if extraction is None or not retry_on_null:
            return extraction

        for attempt in range(self.max_retries):
            if not extraction.has_missing_fields():
                break
            logger.info(
                "Missing fields detected – retrying (attempt %d/%d).",
                attempt + 1,
                self.max_retries,
            )
            second = self._call_llm(input_text)
            if second is not None:
                extraction = self.schema_class.merge(extraction, second)

        return extraction

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _call_llm(self, input_text: str) -> NERBaseModel | None:
        """Build the prompt, call the LLM, and parse the response."""
        prompt = self._prompt_builder.build(input_text)
        raw = self._llm_client.generate(prompt)
        if raw is None:
            return None
        return self.schema_class.safe_parse(raw, input_text=input_text)
