"""Named‑Entity Recognition (NER) predictive task."""

from __future__ import annotations

from collections.abc import Iterable, Sequence
from pathlib import Path
from typing import Any, override

import datasets
import pydantic

from sieves.data import Doc
from sieves.engines import (
    EngineType,
    dspy_,
    glix_,
    huggingface_,
    instructor_,
    langchain_,
    ollama_,
    outlines_,
    vllm_,
)
from sieves.engines.types import GenerationSettings
from sieves.serialization import Config
from sieves.tasks.postprocessing.distillation.types import DistillationFramework
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.ner.bridges import (
    VLLMNER,
    DSPyNER,
    GliXNER,
    InstructorNER,
    LangChainNER,
    OllamaNER,
    OutlinesNER,
)

_TaskModel = (
    dspy_.Model | glix_.Model | instructor_.Model | langchain_.Model | ollama_.Model | outlines_.Model | vllm_.Model
)
_TaskPromptSignature = Any
_TaskResult = (
    list[tuple[str, int, int]]
    | list[tuple[str, str, int, int]]
    | pydantic.BaseModel
    | dspy_.Result
    | glix_.Result
    | huggingface_.Result
    | instructor_.Result
    | langchain_.Result
    | ollama_.Result
    | outlines_.Result
    | vllm_.Result
)
_TaskBridge = DSPyNER | GliXNER | InstructorNER | LangChainNER | OllamaNER | OutlinesNER | VLLMNER


class Entity(pydantic.BaseModel):
    """Entity mention with text span and type."""

    text: str
    context: str
    entity_type: str


class TaskFewshotExample(pydantic.BaseModel):
    """Few‑shot example with entities annotated in text."""

    text: str
    entities: list[Entity]


class NER(PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]):
    """Extract named entities from text using various engines."""

    def __init__(
        self,
        entities: list[str],
        model: _TaskModel,
        task_id: str | None = None,
        include_meta: bool = True,
        batch_size: int = -1,
        prompt_template: str | None = None,
        prompt_signature_desc: str | None = None,
        fewshot_examples: Iterable[TaskFewshotExample] = (),
        generation_settings: GenerationSettings = GenerationSettings(),
    ) -> None:
        """
        Initialize NER task.

        :param entities: List of entities to extract.
        :param model: Model to use.
        :param task_id: Task ID.
        :param include_meta: Whether to include meta information generated by the task.
        :param batch_size: Batch size to use for inference. Use -1 to process all documents at once.
        :param prompt_template: Custom prompt template. If None, task's default template is being used.
        :param prompt_signature_desc: Custom prompt signature description. If None, default will be used.
        :param fewshot_examples: Few-shot examples.
        :param generation_settings: Settings for structured generation.
        """
        self._entities = entities or ["PERSON", "LOCATION", "ORGANIZATION"]
        super().__init__(
            model=model,
            task_id=task_id,
            include_meta=include_meta,
            batch_size=batch_size,
            overwrite=False,
            prompt_template=prompt_template,
            prompt_signature_desc=prompt_signature_desc,
            fewshot_examples=fewshot_examples,
            generation_settings=generation_settings,
        )
        self._fewshot_examples: Iterable[TaskFewshotExample]

    @override
    def _init_bridge(self, engine_type: EngineType) -> _TaskBridge:
        bridge_types = {
            EngineType.langchain: LangChainNER,
            EngineType.ollama: OllamaNER,
            EngineType.outlines: OutlinesNER,
            EngineType.dspy: DSPyNER,
            EngineType.instructor: InstructorNER,
            EngineType.glix: GliXNER,
            EngineType.vllm: VLLMNER,
        }
        try:
            bridge_class = bridge_types[engine_type]
            result = bridge_class(
                task_id=self._task_id,
                prompt_template=self._custom_prompt_template,
                prompt_signature_desc=self._custom_prompt_signature_desc,
                entities=self._entities,
            )
            return result  # type: ignore[return-value]
        except KeyError as err:
            raise KeyError(f"Engine type {engine_type} is not supported by {self.__class__.__name__}.") from err

    @override
    @property
    def supports(self) -> set[EngineType]:
        return {
            EngineType.langchain,
            EngineType.ollama,
            EngineType.dspy,
            EngineType.outlines,
            EngineType.instructor,
            EngineType.glix,
            EngineType.vllm,
        }

    @override
    def _validate_fewshot_examples(self) -> None:
        for fs_example in self._fewshot_examples or []:
            for entity in fs_example.entities:
                if entity.entity_type not in self._entities:
                    raise ValueError(f"Entity {entity.entity_type} not in {self._entities}.")

    @override
    @property
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "entities": self._entities,
        }

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float = 0.5) -> datasets.Dataset:
        # Define metadata and features for the dataset
        features = datasets.Features(
            {
                "text": datasets.Value("string"),
                "entities": datasets.Sequence(
                    datasets.Features(
                        {
                            "text": datasets.Value("string"),
                            "start": datasets.Value("int32"),
                            "end": datasets.Value("int32"),
                            "entity_type": datasets.Value("string"),
                        }
                    )
                ),
            }
        )

        info = datasets.DatasetInfo(
            description=f"Named Entity Recognition dataset with entity types {self._entities}. Generated with sieves "
            f"v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset
        try:
            data: list[tuple[str, list[dict[str, Any]]]] = []
            for doc in docs:
                if self._task_id not in doc.results:
                    raise KeyError(f"Document does not have results for task ID {self._task_id}")

                # Get the entities from the document results
                result = doc.results[self._task_id].entities
                entities: list[dict[str, Any]] = []

                # List format (could be list of dictionaries or other entities)
                for entity in result:
                    assert hasattr(entity, "text")
                    assert hasattr(entity, "start")
                    assert hasattr(entity, "end")
                    assert hasattr(entity, "entity_type")

                    entities.append(
                        {
                            "text": entity.text,
                            "start": entity.start,
                            "end": entity.end,
                            "entity_type": entity.entity_type,
                        }
                    )

                data.append((doc.text or "", entities))

        except KeyError as err:
            raise KeyError(f"Not all documents have results for this task with ID {self._task_id}") from err

        def generate_data() -> Iterable[dict[str, Any]]:
            """Yield results as dicts.

            :return: Results as dicts.
            """
            for text, entities in data:
                yield {"text": text, "entities": entities}

        # Create dataset
        return datasets.Dataset.from_generator(generate_data, features=features, info=info)

    @override
    def distill(
        self,
        base_model_id: str,
        framework: DistillationFramework,
        data: datasets.Dataset | Sequence[Doc],
        output_path: Path | str,
        val_frac: float,
        init_kwargs: dict[str, Any] | None = None,
        train_kwargs: dict[str, Any] | None = None,
        seed: int | None = None,
    ) -> None:
        raise NotImplementedError
