"""Crime-report extraction schema – standalone example.

Demonstrates how to build a domain-specific NER schema on top of llm-ner.
Copy and adapt this file to define your own extraction targets.
"""

from __future__ import annotations

from pydantic import Field

from llmner import EvidenceField, NERBaseModel, NERExtractor, SchemaRegistry

# ---------------------------------------------------------------------------
# 1. Create a registry – one per schema
# ---------------------------------------------------------------------------

registry = SchemaRegistry()

# ---------------------------------------------------------------------------
# 2. Define field types using the registry factories
# ---------------------------------------------------------------------------

GenderType = registry.categorical(
    "gender",
    options=["male", "female"],
    instruction="Extract the subject's gender (infer from pronouns if needed).",
)

NationalityType = registry.categorical(
    "nationality",
    options=[
        "spanish",
        "moroccan",
        "romanian",
        "colombian",
        "ecuadorian",
        "dominican",
        "venezuelan",
        "peruvian",
        "bolivian",
        "argentinian",
        "french",
        "italian",
        "british",
        "german",
        "american",
        "senegalese",
        "algerian",
        "chinese",
        "indian",
        "pakistani",
        "portuguese",
        "russian",
        "ukrainian",
        "bulgarian",
        "south american",
        "latin american",
        "african",
        "asian",
        "european",
    ],
    instruction=(
        "Extract the subject's nationality or ethnicity. "
        "Always use singular masculine form."
    ),
)

AgeType = registry.int_range(
    "age",
    "Extract the subject's age as an integer or range (e.g. '25' or '25-30').",
)

HeightType = registry.int_range(
    "height",
    (
        "Extract the subject's height in whole centimetres or range "
        "(e.g. '175' or '170-180'). Convert metres to centimetres."
    ),
)

HairType = registry.categorical(
    "hair_colour",
    options=["black", "dark brown", "brown", "blonde", "red", "grey", "bald"],
    instruction=(
        "Extract hair colour, ignoring cut or style. "
        "Map synonyms to the nearest allowed value."
    ),
)

WeaponType = registry.categorical(
    "weapon_method",
    options={
        "physical_force": "body used to attack or restrain (punches, pushes, grabs)",
        "bladed_weapon": "cutting or stabbing objects (knives, blades)",
        "blunt_weapon": "blunt objects used to strike (sticks, metal bars, stones)",
        "firearm": "pistols, revolvers, or similar",
        "verbal_threat": "verbal threats without any other weapon",
    },
    instruction=(
        "Extract the weapon or method used. Map synonyms to the nearest allowed value."
    ),
    allow_multiple=True,
)

DatetimeType = registry.datetime_format(
    "datetime",
    "Extract the date and time of the incident. Strict format: YYYY-MM-DD HH:MM:SS.",
)

AddressType = registry.generic(
    "address",
    "Extract the full address where the incident occurred.",
)

# ---------------------------------------------------------------------------
# 3. Define the Pydantic models
# ---------------------------------------------------------------------------


class IncidentFacts(NERBaseModel):
    """Factual circumstances of the reported incident."""

    datetime: DatetimeType | None = Field(None, alias="datetime")  # type: ignore[valid-type]
    address: AddressType | None = Field(None, alias="address")  # type: ignore[valid-type]


class IncidentMO(NERBaseModel):
    """Modus operandi of the incident."""

    weapon_method: WeaponType | None = Field(None, alias="weapon_method")  # type: ignore[valid-type]


class Suspect(NERBaseModel):
    """Physical description of one perceived suspect."""

    gender: GenderType | None = Field(None, alias="gender")  # type: ignore[valid-type]
    nationality: NationalityType | None = Field(None, alias="nationality")  # type: ignore[valid-type]
    age: AgeType | None = Field(None, alias="age")  # type: ignore[valid-type]
    height: HeightType | None = Field(None, alias="height")  # type: ignore[valid-type]
    hair_colour: HairType | None = Field(None, alias="hair_colour")  # type: ignore[valid-type]

    def summary(self) -> str:
        """Return a human-readable summary of this suspect's description."""

        def _v(field: EvidenceField | None) -> str:
            return field.value if field and field.value else "unknown"

        return (
            f"{_v(self.gender)}, {_v(self.nationality)}, "
            f"age {_v(self.age)}, height {_v(self.height)} cm, "
            f"hair {_v(self.hair_colour)}"
        )


class CrimeReport(NERBaseModel):
    """Root extraction schema for a crime / incident report."""

    facts: IncidentFacts | None = Field(None, alias="facts")
    modus_operandi: IncidentMO | None = Field(None, alias="modus_operandi")
    suspects: list[Suspect] = Field(default_factory=list, alias="suspects")


# ---------------------------------------------------------------------------
# 4. Extractor factory
# ---------------------------------------------------------------------------

SYSTEM_ROLE = (
    "You are an expert forensic analyst specialised in extracting "
    "structured data from police incident reports."
)

SYSTEM_TASK = (
    "Extract objective information from the incident report text. "
    "If a piece of data is NOT present in the text, set its 'value' to null "
    "and its 'evidence' to null. "
    "DO NOT invent or infer data that is not explicitly stated. "
    "For 'evidence', extract a short verbatim quote (ideally 1-8 words) that "
    "unambiguously justifies the extracted value. "
    "If multiple suspects are described, create one entry per suspect. "
    "Return ONLY a valid JSON object matching the provided schema."
)


def create_extractor(
    llm_base_url: str = "http://localhost:11434",
    llm_model: str = "qwen2.5:7b-instruct",
    llm_timeout: int = 120,
    max_retries: int = 1,
) -> NERExtractor:
    """Return a configured :class:`~llmner.extractor.NERExtractor`.

    Args:
        llm_base_url: Ollama server URL.
        llm_model:    Ollama model tag.
        llm_timeout:  Request timeout in seconds.
        max_retries:  Retry attempts on incomplete extraction.
    """
    return NERExtractor(
        schema_class=CrimeReport,
        system_role=SYSTEM_ROLE,
        system_task=SYSTEM_TASK,
        rules_registry=registry.rules,
        llm_base_url=llm_base_url,
        llm_model=llm_model,
        llm_timeout=llm_timeout,
        max_retries=max_retries,
    )
