"""Crime-report NER schema used by the accuracy integration tests.

Defines:

* :data:`registry`            – :class:`~llmner.factories.SchemaRegistry`
  instance with all crime-specific field types.
* :class:`CrimeFacts`         – datetime, address, number of suspects.
* :class:`CrimeMO`            – modus operandi (weapon / method).
* :class:`CrimeSuspect`       – physical description of one suspect.
* :class:`CrimeExtraction`    – root extraction model.
* :func:`create_crime_extractor` – factory that returns a ready-to-use
  :class:`~llmner.extractor.NERExtractor`.

The prompts are in Spanish because the target corpus is composed of Spanish
police complaint reports.
"""

from __future__ import annotations

from pydantic import Field

from llmner import EvidenceField, NERBaseModel, NERExtractor, SchemaRegistry

# ---------------------------------------------------------------------------
# Registry and field types
# ---------------------------------------------------------------------------

registry = SchemaRegistry()

GenderType = registry.categorical(
    "género",
    options=["masculino", "femenino"],
    instruction=(
        "Extrae el género del sujeto. "
        "Infiere a partir de pronombres, sustantivos o adjetivos "
        "(ej. 'mujer'/'señora'/'la autora' → femenino, "
        "'varón'/'hombre'/'el autor' → masculino). "
        "El value DEBE ser 'masculino' o 'femenino', nunca 'mujer'/'hombre'."
    ),
)

ComplexionType = registry.categorical(
    "complexión",
    options={
        "delgado": "poco peso o masa muscular",
        "normal": "complexión media o estándar",
        "atletico": "musculoso o en buena forma física",
        "fuerte": "corpulento o robusto",
        "sobrepeso_obeso": "exceso de peso",
    },
    instruction=(
        "Extrae la complexión física. Mapea sinónimos al valor permitido más cercano "
    ),
)

HairType = registry.categorical(
    "color_pelo",
    options=[
        "negro",
        "castaño_oscuro",
        "castaño",
        "rubio",
        "pelirrojo",
        "canoso",
        "calvo",
    ],
    instruction=(
        "Extrae el color de pelo EXACTO que aparece en el texto, omitiendo cortes o estilos. "
        "NO añadas matices que no estén explícitamente en el texto "
        "Mapea sinónimos solo cuando sea necesario (ej. moreno → castaño_oscuro)."
    ),
)

EyesType = registry.categorical(
    "color_ojos",
    options=["oscuros", "marrones", "avellana", "verdes", "azules", "claros"],
    instruction=(
        "Extrae el color de ojos. "
        "Mapea sinónimos al valor permitido más cercano "
        "Mapea sinónimos solo cuando sea necesario (ej. miel → avellana)."
    ),
)

WeaponMoType = registry.categorical(
    "mo_arma",
    options={
        "fuerza_física": (
            "uso del cuerpo para agredir o someter (ej. puñetazos, empujones, tirones)"
        ),
        "arma_blanca": (
            "uso de objetos cortantes o punzantes (ej. navajas, cuchillos)"
        ),
        "arma_contundente": (
            "uso de objetos romos para golpear (ej. palos, barras de metal, piedras)"
        ),
        "arma_fuego": "uso de pistolas, revólveres o similares",
        "intimidación_verbal": ("amenazas verbales sin uso o amago de otras armas"),
    },
    instruction=(
        "Extrae el tipo de arma o método utilizado en el modus operandi. "
        "Mapea sinónimos al valor permitido más cercano."
    ),
    allow_multiple=True,
)

_NATIONALITIES = [
    "español",
    "marroquí",
    "rumano",
    "colombiano",
    "ecuatoriano",
    "dominicano",
    "venezolano",
    "peruano",
    "boliviano",
    "argentino",
    "brasileño",
    "paraguayo",
    "uruguayo",
    "chileno",
    "mexicano",
    "cubano",
    "estadounidense",
    "francés",
    "italiano",
    "británico",
    "alemán",
    "senegalés",
    "argelino",
    "chino",
    "indio",
    "pakistaní",
    "portugués",
    "ruso",
    "ucraniano",
    "búlgaro",
    "árabe",
    "sudamericano",
    "latinoamericano",
    "gitano",
    "africano",
    "asiático",
    "europeo",
]

_NATIONALITY_FEM_TO_MASC: dict[str, str] = {
    "española": "español",
    "rumana": "rumano",
    "colombiana": "colombiano",
    "ecuatoriana": "ecuatoriano",
    "dominicana": "dominicano",
    "venezolana": "venezolano",
    "peruana": "peruano",
    "boliviana": "boliviano",
    "argentina": "argentino",
    "brasileña": "brasileño",
    "paraguaya": "paraguayo",
    "uruguaya": "uruguayo",
    "chilena": "chileno",
    "mexicana": "mexicano",
    "cubana": "cubano",
    "francesa": "francés",
    "italiana": "italiano",
    "británica": "británico",
    "alemana": "alemán",
    "senegalesa": "senegalés",
    "argelina": "argelino",
    "china": "chino",
    "india": "indio",
    "portuguesa": "portugués",
    "rusa": "ruso",
    "ucraniana": "ucraniano",
    "búlgara": "búlgaro",
    "sudamericana": "sudamericano",
    "latinoamericana": "latinoamericano",
    "gitana": "gitano",
    "africana": "africano",
    "asiática": "asiático",
    "europea": "europeo",
}

NationalityType = registry.categorical(
    "nacionalidad",
    options=_NATIONALITIES,
    instruction=(
        "Extrae la nacionalidad o etnia del sujeto. "
        "Exprésalo SIEMPRE en español y en masculino singular "
        "(ej. 'española' → 'español', 'rumana' → 'rumano', "
        "'spanish' → 'español', 'moroccan' → 'marroquí'). "
    ),
    replacements=_NATIONALITY_FEM_TO_MASC,
)

AgeType = registry.int_range(
    "edad",
    "Extrae la edad del sujeto. Formato: número entero o rango (ej. '25-30').",
)

HeightType = registry.int_range(
    "altura",
    (
        "Extrae la altura del sujeto en centímetros enteros o rango "
        "(ej. '175', '170-180'). Convierte metros a centímetros. "
        "IMPORTANTE: si el texto indica un rango (ej. 'entre 1,65 y 1,80 metros'), "
        "devuelve el rango completo '165-180', NUNCA la media ni un solo extremo."
    ),
)

DatetimeType = registry.datetime_format(
    "datetime_perc",
    (
        "Extrae la fecha y hora EXACTA en que ocurrieron los hechos "
        "(NO la hora de comparecencia). "
        "Busca expresiones como 'sobre las HH:MM horas del día YYYY-MM-DD'. "
        "Formato estricto de salida: YYYY-MM-DD HH:MM:SS. "
        "NUNCA devuelvas null si el texto contiene fecha y hora de los hechos."
    ),
)

AddressType = registry.generic(
    "address_perc",
    "Extrae la dirección completa donde ocurrieron los hechos.",
)

NumPerceivedType = registry.int_range(
    "num_perceived_suspects",
    (
        "Extrae el número total de sospechosos percibidos. "
        "Formato: número entero o rango (ej. '2', '2-3')."
    ),
)

# ---------------------------------------------------------------------------
# Extraction models
# ---------------------------------------------------------------------------


class CrimeFacts(NERBaseModel):
    """Factual circumstances of the reported crime."""

    datetime_perc: DatetimeType | None = Field(None, alias="datetime_perc")  # type: ignore[valid-type]
    address_perc: AddressType | None = Field(None, alias="address_perc")  # type: ignore[valid-type]
    num_perceived_suspects: NumPerceivedType | None = Field(  # type: ignore[valid-type]
        None, alias="num_perceived_suspects"
    )


class CrimeMO(NERBaseModel):
    """Modus operandi of the crime."""

    weapon_mo: WeaponMoType | None = Field(None, alias="mo_arma")  # type: ignore[valid-type]


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

    gender: GenderType | None = Field(None, alias="género")  # type: ignore[valid-type]
    nationality: NationalityType | None = Field(None, alias="nacionalidad")  # type: ignore[valid-type]
    age: AgeType | None = Field(None, alias="edad")  # type: ignore[valid-type]
    height: HeightType | None = Field(None, alias="altura")  # type: ignore[valid-type]
    complexion: ComplexionType | None = Field(None, alias="complexión")  # type: ignore[valid-type]
    hair: HairType | None = Field(None, alias="color_pelo")  # type: ignore[valid-type]
    eyes: EyesType | None = Field(None, alias="color_ojos")  # type: ignore[valid-type]

    def to_row(self) -> dict[str, str | None]:
        """Return a flat dict suitable for comparison with ground-truth CSV rows.

        Keys follow the Spanish alias names used in the ground-truth data
        (``género``, ``nacionalidad``, etc.) plus ``evidence_<alias>``
        companions.
        """
        mapping: list[tuple[str, str]] = [
            ("gender", "género"),
            ("nationality", "nacionalidad"),
            ("age", "edad"),
            ("height", "altura"),
            ("complexion", "complexión"),
            ("hair", "color_pelo"),
            ("eyes", "color_ojos"),
        ]
        row: dict[str, str | None] = {}
        for field_name, alias in mapping:
            field_val: EvidenceField | None = getattr(self, field_name)
            row[alias] = field_val.value if field_val else None
            row[f"evidence_{alias}"] = field_val.evidence if field_val else None
        return row


class CrimeExtraction(NERBaseModel):
    """Root extraction schema for a Spanish police complaint report."""

    facts: CrimeFacts | None = Field(None, alias="facts")
    modus_operandi: CrimeMO | None = Field(None, alias="modus_operandi")
    suspects: list[CrimeSuspect] = Field(default_factory=list, alias="suspects")


# ---------------------------------------------------------------------------
# LLM prompt configuration
# ---------------------------------------------------------------------------

SYSTEM_ROLE = (
    "Eres un analista experto en inteligencia policial y extracción forense "
    "de datos estructurados a partir de denuncias policiales españolas."
)

SYSTEM_TASK = (
    "Extrae información objetiva del texto de la denuncia. "
    "Si un dato NO aparece en el texto, pon null en su campo 'value' y null en 'evidence'. "
    "NO inventes ni asumas datos que no estén en el texto. "
    "Para el campo 'evidence', extrae una cita literal corta que JUSTIFIQUE INEQUÍVOCAMENTE el 'value'. "
    "La cita elegida debe ser suficiente por sí misma "
    "(ej. debe incluir fecha Y hora, o usar el vocablo inequívoco como 'varón' "
    "en lugar del ambiguo 'individuo'). "
    "Mantenlo corto (ideal 1-8 palabras) pero prioriza que la evidencia sea "
    "irrefutable y completa. EVITA oraciones enteras. "
    "REGLAS ADICIONALES: "
    "1. Todos los 'value' DEBEN estar en ESPAÑOL. NUNCA en inglés ni otro idioma. "
    "2. Si la denuncia describe a varios sospechosos, crea una entrada por cada uno. "
    "3. Las claves del JSON generado deben coincidir EXACTAMENTE con las del "
    "esquema proporcionado. Todos los atributos terminales ahora esperan un "
    "objeto con 'value' y 'evidence'. "
    "4. Devuelve ÚNICAMENTE un objeto JSON válido acorde al esquema. "
    "Sin texto adicional. "
    "5. Lee el texto COMPLETO antes de extraer. No dejes campos en null si la "
    "información está presente en cualquier parte del texto. "
    "6. Para rangos numéricos (edad, altura), conserva SIEMPRE el rango original "
    "del texto (ej. 'entre 25 y 30' → '25-30'). NUNCA promedies ni redondees."
)

PROMPT_TEMPLATE = """\
INSTRUCCIÓN DE EXTRACCIÓN DE ENTIDADES

ROL: {system_role}

TAREA: {system_task}

FORMATO ESPECÍFICO DE CAMPOS:
{rules_text}

ESQUEMA DE EXTRACCIÓN A REPLICAR ESTRICTAMENTE (USA ESTAS MISMAS CLAVES):
{schema_json}

TEXTO A PROCESAR:
"{input_text}"
"""


# ---------------------------------------------------------------------------
# Extractor factory
# ---------------------------------------------------------------------------


def create_crime_extractor(
    llm_base_url: str = "http://localhost:11434",
    llm_model: str = "qwen2.5:7b-instruct",
    llm_timeout: int = 120,
    max_retries: int = 2,
) -> NERExtractor:
    """Return a :class:`~llmner.extractor.NERExtractor` configured for
    Spanish police complaint reports.

    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.

    Returns:
        A ready-to-use :class:`~llmner.extractor.NERExtractor`.
    """
    return NERExtractor(
        schema_class=CrimeExtraction,
        system_role=SYSTEM_ROLE,
        system_task=SYSTEM_TASK,
        rules_registry=registry.rules,
        prompt_template=PROMPT_TEMPLATE,
        llm_base_url=llm_base_url,
        llm_model=llm_model,
        llm_timeout=llm_timeout,
        max_retries=max_retries,
    )
