"""Core task implementation."""

from __future__ import annotations

import abc
import itertools
import sys
from collections.abc import Callable, Iterable, Iterator
from typing import TYPE_CHECKING, Any

import dspy

from sieves.data import Doc
from sieves.serialization import Attribute, Config
from sieves.tasks.predictive.evaluation import TaskEvaluationReport

if TYPE_CHECKING:
    from sieves.pipeline import Pipeline


class Task(abc.ABC):
    """Abstract base class for tasks that can be executed on documents."""

    def __init__(
        self,
        task_id: str | None,
        include_meta: bool,
        batch_size: int,
        condition: Callable[[Doc], bool] | None = None,
    ):
        """
        Initiate new Task.

        :param task_id: Task ID.
        :param include_meta: Whether to include meta information generated by the task.
        :param batch_size: Batch size for processing documents. Use -1 to process all documents at once.
        :param condition: Optional callable that determines whether to process each document.
                          If provided, called with each Doc; if returns False, document is skipped
                          and results[task_id] is set to None.
        """
        self._task_id = task_id if task_id else self.__class__.__name__
        self._include_meta = include_meta
        self._batch_size = batch_size
        self._condition = condition

    @property
    def id(self) -> str:
        """Return task ID.

        Used by pipeline for results and dependency management.

        :return: Task ID.
        """
        return self._task_id

    def __call__(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Execute task with conditional logic.

        Checks the condition for each document without materializing all docs upfront.
        Passes all documents that pass the condition to _call() for proper batching.
        Documents that fail the condition have results[task_id] set to None.

        :param docs: Docs to process.
        :return: Processed docs (in original order).
        """
        docs = iter(docs) if not isinstance(docs, Iterator) else docs

        # Materialize docs in batches. This doesn't incur additional memory overhead, as docs are materialized in
        # batches downstream anyway.
        batch_size = self._batch_size if self._batch_size > 0 else sys.maxsize
        while docs_batch := [doc for doc in itertools.islice(docs, batch_size)]:
            # First pass: determine which docs pass the condition by index.
            passing_indices: set[int] = {
                idx for idx, doc in enumerate(docs_batch) if self._condition is None or self._condition(doc)
            }

            # Process all passing docs in one batch.
            processed = self._call(d for i, d in enumerate(docs_batch) if i in passing_indices)
            processed_iter = iter(processed) if not isinstance(processed, Iterator) else processed

            # Iterate through original docs in order and yield results.
            for idx, doc in enumerate(docs_batch):
                if idx in passing_indices:
                    # Doc passed condition - use processed result.
                    yield next(processed_iter)
                else:
                    # Doc failed condition - set `None` result and yield original.
                    doc.results[self.id] = None
                    yield doc

    @abc.abstractmethod
    def _call(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Execute task logic (to be implemented by subclasses).

        :param docs: Docs to process.
        :return: Processed docs.
        """

    def evaluate(self, docs: Iterable[Doc], judge: dspy.LM | None = None) -> TaskEvaluationReport:
        """Evaluate task performance.

        :param docs: Documents to evaluate.
        :param judge: Optional judge model for evaluation.
        :return: Evaluation report.
        """
        raise NotImplementedError(f"Evaluation not implemented for task {self.__class__.__name__}")

    def __add__(self, other: Task | Pipeline) -> Pipeline:
        """Chain this task with another task or pipeline using the ``+`` operator.

        This returns a new ``Pipeline`` that executes this task first, followed by the
        task(s) in ``other``. The original task(s)/pipeline are not mutated.

        Cache semantics:
        - If ``other`` is a ``Pipeline``, the resulting pipeline adopts ``other``'s
          ``use_cache`` setting (because the left-hand side is a single task).
        - If ``other`` is a ``Task``, the resulting pipeline defaults to ``use_cache=True``.

        :param other: A ``Task`` or ``Pipeline`` to execute after this task.
        :return: A new ``Pipeline`` representing the chained execution.
        :raises TypeError: If ``other`` is not a ``Task`` or ``Pipeline``.
        """
        # Lazy import to avoid circular dependency at module import time.
        from sieves.pipeline import Pipeline

        if isinstance(other, Pipeline):
            return Pipeline(tasks=[self, *other.tasks], use_cache=other.use_cache)

        if isinstance(other, Task):
            return Pipeline(tasks=[self, other])

        raise TypeError(f"Cannot chain Task with {type(other).__name__}")

    @property
    def _state(self) -> dict[str, Any]:
        """Return attributes to serialize.

        :return: Dict of attributes to serialize.
        """
        return {
            "task_id": self._task_id,
            "include_meta": self._include_meta,
            "batch_size": self._batch_size,
            "condition": self._condition,
        }

    def serialize(self) -> Config:
        """Serialize task.

        :return: Config instance.
        """
        return Config.create(self.__class__, {k: Attribute(value=v) for k, v in self._state.items()})

    @classmethod
    def deserialize(cls, config: Config, **kwargs: dict[str, Any]) -> Task:
        """Generate Task instance from config.

        :param config: Config to generate instance from.
        :param kwargs: Values to inject into loaded config.
        :return: Deserialized Task instance.
        """
        # Deserialize and inject model wrapper.
        return cls(**config.to_init_dict(cls, **kwargs))
