"""Feedback types for the optimization system.

This module provides the FeedbackType enum and Feedback dataclass which
represent evaluation results from loss functions. Feedback objects can
propagate backward through the computation graph to accumulate improvement
suggestions into Parameters.

The API mirrors PyTorch's loss.backward() pattern:
    >>> feedback = await loss_fn(output, target, record=record)
    >>> await feedback.backward()  # Propagates feedback to Parameters

Example:
    >>> from plait.optimization.feedback import Feedback, FeedbackType
    >>>
    >>> # Create feedback from evaluation
    >>> feedback = Feedback(
    ...     content="Response was too verbose",
    ...     score=0.6,
    ...     feedback_type=FeedbackType.LLM_JUDGE,
    ... )
    >>> str(feedback)
    '[0.60] Response was too verbose'
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from plait.optimization.record import ForwardRecord


class FeedbackType(Enum):
    """Classification of feedback sources.

    Different feedback types may be weighted or processed differently
    during optimization. This enum allows tracking the provenance of
    feedback throughout the backward pass.

    Attributes:
        HUMAN: Feedback provided directly by a human evaluator.
        LLM_JUDGE: Feedback generated by an LLM acting as a judge.
        VERIFIER: Feedback from programmatic verification (code-based checks).
        COMPOSITE: Aggregated feedback from multiple sources.

    Example:
        >>> feedback = Feedback(
        ...     content="Good response",
        ...     feedback_type=FeedbackType.HUMAN,
        ... )
        >>> feedback.feedback_type == FeedbackType.HUMAN
        True
    """

    HUMAN = "human"
    LLM_JUDGE = "llm_judge"
    VERIFIER = "verifier"
    COMPOSITE = "composite"


@dataclass
class Feedback:
    """Feedback on a module output, analogous to a loss tensor in PyTorch.

    Holds the evaluation result and (optionally) a reference to the forward
    pass record for backward propagation. When `backward()` is called,
    feedback propagates through the computation graph to accumulate
    improvement suggestions into Parameters.

    Attributes:
        content: The feedback text describing what was good or bad about
            the output and how it could be improved.
        score: Optional numeric score normalized to 0-1 range, where 1.0
            indicates perfect output and 0.0 indicates complete failure.
        feedback_type: Classification of the feedback source. Defaults
            to HUMAN if not specified.
        metadata: Optional dictionary for storing additional information
            such as raw scores, criteria used, or evaluation context.

    Example:
        >>> # Create feedback with score
        >>> feedback = Feedback(
        ...     content="Response was helpful but too verbose",
        ...     score=0.7,
        ...     feedback_type=FeedbackType.LLM_JUDGE,
        ...     metadata={"criteria": "helpfulness"},
        ... )
        >>> str(feedback)
        '[0.70] Response was helpful but too verbose'
        >>>
        >>> # Create feedback without score
        >>> feedback = Feedback(content="Add more examples")
        >>> str(feedback)
        'Add more examples'
        >>>
        >>> # backward() requires a ForwardRecord
        >>> import asyncio
        >>> asyncio.run(feedback.backward())  # Raises RuntimeError

    Note:
        The `_records` and `_optimizer` fields are internal and should not
        be set directly. They are populated by loss functions when computing
        feedback. For batch inputs, all records from the batch are attached.
    """

    content: str
    score: float | None = None
    feedback_type: FeedbackType = FeedbackType.HUMAN
    metadata: dict[str, Any] = field(default_factory=dict)

    # Internal: references to forward passes for backward()
    # For aggregated batch loss, holds records from all samples in the batch
    _records: list[ForwardRecord] = field(
        default_factory=list, repr=False, compare=False
    )
    _optimizer: Any = field(
        default=None, repr=False, compare=False
    )  # Optimizer type hint avoided for circular import

    async def backward(self, optimizer: Any = None) -> None:
        """Propagate feedback backward through the computation graph.

        This is the primary API for backward passes, mirroring loss.backward()
        in PyTorch. Feedback accumulates into Parameter._feedback_buffer for
        later aggregation by the optimizer.

        For aggregated batch loss (when Loss is called with a list of outputs),
        this method iterates through all attached records and propagates
        feedback to each one concurrently. This matches PyTorch semantics
        where loss.backward() on a reduced loss propagates gradients to all
        batch samples.

        Args:
            optimizer: Optional optimizer providing reasoning LLM for
                custom backward implementations. If not provided, uses
                self._optimizer if set.

        Raises:
            RuntimeError: If no ForwardRecords are attached. This happens when
                feedback was computed without training mode or without passing
                TracedOutput to the loss function.

        Example:
            >>> # Single sample
            >>> module.train()
            >>> output = await module("Hello")
            >>> feedback = await loss_fn(output, target)
            >>> await feedback.backward()  # Propagates to one record
            >>>
            >>> # Batch training (aggregated loss)
            >>> outputs = [await module(inp) for inp in batch]
            >>> feedback = await loss_fn(outputs, targets)  # Single aggregated feedback
            >>> await feedback.backward()  # Propagates to all records

        Note:
            The actual backward propagation is implemented in the backward
            module's `_propagate_backward()` function. This method validates
            records are present and then delegates to that function for each.
        """
        if not self._records:
            raise RuntimeError(
                "Cannot call backward() without ForwardRecords. "
                "Use training mode (module.train()) or pass TracedOutput to loss."
            )

        # Import here to avoid circular imports
        import asyncio

        from plait.optimization.backward import _propagate_backward

        opt = optimizer or self._optimizer
        reasoning_llm = (
            opt.reasoning_llm if opt and hasattr(opt, "reasoning_llm") else None
        )

        # Propagate feedback through all attached records concurrently
        async def propagate_one(record: ForwardRecord) -> None:
            await _propagate_backward(
                feedback=self,
                record=record,
                reasoning_llm=reasoning_llm,
                optimizer=opt,
            )

        await asyncio.gather(*[propagate_one(rec) for rec in self._records])

    def __str__(self) -> str:
        """Return string representation of the feedback.

        If a score is present, it's prepended to the content in brackets.

        Returns:
            Formatted feedback string.

        Example:
            >>> str(Feedback(content="Good job", score=0.9))
            '[0.90] Good job'
            >>> str(Feedback(content="Needs work"))
            'Needs work'
        """
        if self.score is not None:
            return f"[{self.score:.2f}] {self.content}"
        return self.content
