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

from __future__ import annotations

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.model_wrappers import ModelType
from sieves.model_wrappers.types import ModelSettings
from sieves.serialization import Config
from sieves.tasks.distillation.types import DistillationFramework
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.pii_masking.bridges import DSPyPIIMasking, PydanticPIIMasking
from sieves.tasks.predictive.schemas.pii_masking import (
    FewshotExample,
    TaskModel,
    TaskPromptSignature,
    TaskResult,
)

_TaskBridge = DSPyPIIMasking | PydanticPIIMasking


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

    def __init__(
        self,
        model: TaskModel,
        pii_types: Sequence[str] | dict[str, str] | None = None,
        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] = (),
        model_settings: ModelSettings = ModelSettings(),
        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 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 model_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 = "[MASKED]"

        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,
            model_settings=model_settings,
            condition=condition,
        )

    @property
    @override
    def fewshot_example_type(self) -> type[FewshotExample]:
        """Return few-shot example type.

        :return: Few-shot example type.
        """
        return FewshotExample

    @property
    @override
    def prompt_signature(self) -> type[pydantic.BaseModel]:
        return TaskResult

    @property
    @override
    def metric(self) -> str:
        return "F1"

    @override
    def _compute_metrics(self, truths: list[Any], preds: list[Any], judge: dspy.LM | None = None) -> dict[str, float]:
        """Compute corpus-level metrics.

        :param truths: List of ground truths.
        :param preds: List of predictions.
        :param judge: Optional DSPy LM instance to use as judge for generative tasks.
        :return: Dictionary of metrics.
        """
        tp = 0
        fp = 0
        fn = 0

        for gold, pred in zip(truths, preds):
            # Extract entities.
            if gold is not None:
                assert isinstance(gold, TaskResult)
                true_entities = {(e.entity_type, e.text) for e in gold.pii_entities}
            else:
                true_entities = set()

            if pred is not None:
                assert isinstance(pred, TaskResult)
                pred_entities = {(e.entity_type, e.text) for e in pred.pii_entities}
            else:
                pred_entities = set()

            tp += len(true_entities & pred_entities)
            fp += len(pred_entities - true_entities)
            fn += len(true_entities - pred_entities)

        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0

        return {self.metric: f1}

    @override
    def _init_bridge(self, model_type: ModelType) -> _TaskBridge:
        bridge_types: dict[ModelType, type[_TaskBridge]] = {
            ModelType.dspy: DSPyPIIMasking,
            ModelType.langchain: PydanticPIIMasking,
            ModelType.outlines: PydanticPIIMasking,
        }

        try:
            return bridge_types[model_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,
                model_settings=self._model_settings,
                prompt_signature=self.prompt_signature,
                model_type=model_type,
                fewshot_examples=self._fewshot_examples,
            )
        except KeyError as err:
            raise KeyError(f"Model type {model_type} is not supported by {self.__class__.__name__}.") from err

    @staticmethod
    @override
    def supports() -> set[ModelType]:
        return {
            ModelType.dspy,
            ModelType.langchain,
            ModelType.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,
        }

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float | None = None) -> 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_dspy_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
