"""Allows masking of PII (Personally Identifiable Information) in text documents."""

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

import datasets
import dspy
import pydantic

from sieves.data.doc import Doc
from sieves.engines import EngineType, dspy_, langchain_, outlines_
from sieves.engines.types import GenerationSettings
from sieves.serialization import Config
from sieves.tasks.distillation.types import DistillationFramework
from sieves.tasks.predictive.core import FewshotExample as BaseFewshotExample
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.pii_masking.bridges import (
    DSPyPIIMasking,
    LangChainPIIMasking,
    OutlinesPIIMasking,
)

_TaskModel = dspy_.Model | langchain_.Model | outlines_.Model
_TaskPromptSignature = pydantic.BaseModel | dspy_.PromptSignature
_TaskResult = pydantic.BaseModel | dspy_.Result
_TaskBridge = DSPyPIIMasking | LangChainPIIMasking | OutlinesPIIMasking


class PIIEntity(pydantic.BaseModel, frozen=True):
    """PII entity."""

    entity_type: str
    text: str


class FewshotExample(BaseFewshotExample):
    """Example for PII masking few-shot prompting."""

    masked_text: str
    pii_entities: list[PIIEntity]

    @override
    @property
    def target_fields(self) -> Sequence[str]:
        return "masked_text", "pii_entities"


class PIIMasking(PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]):
    """Task for masking PII (Personally Identifiable Information) in text documents.

    Examples:
        Default behavior (masks all common PII types):

        >>> masking = PIIMasking(model=model)

        Specify PII types as a list:

        >>> masking = PIIMasking(
        ...     model=model,
        ...     pii_types=["EMAIL", "PHONE", "SSN"],
        ... )

        Using dict format with descriptions for better PII detection:

        >>> masking = PIIMasking(
        ...     model=model,
        ...     pii_types={
        ...         "EMAIL": "Email addresses in any format",
        ...         "PHONE": "Phone numbers including country codes",
        ...         "SSN": "Social Security Numbers (XXX-XX-XXXX format)"
        ...     },
        ... )
    """

    def __init__(
        self,
        model: _TaskModel,
        pii_types: Sequence[str] | dict[str, str] | None = None,
        mask_placeholder: str = "[MASKED]",
        task_id: str | None = None,
        include_meta: bool = True,
        batch_size: int = -1,
        overwrite: bool = False,
        prompt_instructions: str | None = None,
        fewshot_examples: Sequence[FewshotExample] = (),
        generation_settings: GenerationSettings = GenerationSettings(),
        condition: Callable[[Doc], bool] | None = None,
    ) -> None:
        """Initialize PIIMasking task.

        :param model: Model to use.
        :param pii_types: Types of PII to mask. Supports three formats:
            - None (default): Uses all common PII types
            - List format: `["EMAIL", "PHONE", "SSN"]`
            - Dict format: `{"EMAIL": "Email addresses", "PHONE": "Phone numbers"}`
            The dict format allows you to provide descriptions that help the model better identify PII.
        :param mask_placeholder: String to replace PII with.
        :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 overwrite: Whether to overwrite original document text with masked text.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        :param fewshot_examples: Few-shot examples.
        :param generation_settings: Settings for structured generation.
        :param condition: Optional callable that determines whether to process each document.
        """
        if pii_types is not None:
            if isinstance(pii_types, dict):
                self._pii_types = list(pii_types.keys())
                self._pii_type_descriptions = pii_types
            else:
                self._pii_types = list(pii_types)
                self._pii_type_descriptions = {}
        else:
            self._pii_types = None
            self._pii_type_descriptions = {}

        self._pii_types_param = pii_types
        self._mask_placeholder = mask_placeholder

        super().__init__(
            model=model,
            task_id=task_id,
            include_meta=include_meta,
            batch_size=batch_size,
            overwrite=overwrite,
            prompt_instructions=prompt_instructions,
            fewshot_examples=fewshot_examples,
            generation_settings=generation_settings,
            condition=condition,
        )

    @override
    def _init_bridge(self, engine_type: EngineType) -> _TaskBridge:
        bridge_types: dict[EngineType, type[_TaskBridge]] = {
            EngineType.dspy: DSPyPIIMasking,
            EngineType.langchain: LangChainPIIMasking,
            EngineType.outlines: OutlinesPIIMasking,
        }

        try:
            return bridge_types[engine_type](
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                mask_placeholder=self._mask_placeholder,
                pii_types=self._pii_types_param,
                overwrite=self._overwrite,
                generation_settings=self._generation_settings,
            )
        except KeyError as err:
            raise KeyError(f"Engine type {engine_type} is not supported by {self.__class__.__name__}.") from err

    @staticmethod
    @override
    def supports() -> set[EngineType]:
        return {
            EngineType.dspy,
            EngineType.langchain,
            EngineType.outlines,
        }

    @property
    @override
    def _state(self) -> dict[str, Any]:
        # Store pii_types as dict if descriptions exist, else as original value
        pii_types_state = self._pii_type_descriptions if self._pii_type_descriptions else self._pii_types
        return {
            **super()._state,
            "pii_types": pii_types_state,
            "mask_placeholder": self._mask_placeholder,
        }

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float = 0.5) -> datasets.Dataset:
        # Define metadata.
        features = datasets.Features({"text": datasets.Value("string"), "masked_text": datasets.Value("string")})
        info = datasets.DatasetInfo(
            description=f"PII masking dataset. Generated with sieves v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset.
        try:
            data = [(doc.text, doc.results[self._task_id]["masked_text"]) for doc in docs]
        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, masked_text in data:
                yield {"text": text, "masked_text": masked_text}

        # 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

    @override
    def _evaluate_optimization_example(
        self, truth: dspy.Example, pred: dspy.Prediction, trace: Any, model: dspy.LM
    ) -> float:
        # Compute entity detection F1 score based on (entity_type, text) pairs
        true_entities = {(e["entity_type"], e["text"]) for e in truth["pii_entities"]}
        pred_entities = {(e["entity_type"], e["text"]) for e in pred.get("pii_entities", [])}

        if not true_entities:
            return 1.0 if not pred_entities else 0.0

        precision = len(true_entities & pred_entities) / len(pred_entities) if pred_entities else 0
        recall = len(true_entities & pred_entities) / len(true_entities)
        return 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
