"""Base class for ablation strategies."""

from __future__ import annotations

import abc
from dataclasses import dataclass
from typing import Any, Iterator

from obliteratus.models.loader import ModelHandle


@dataclass
class AblationSpec:
    """Describes one atomic ablation operation."""

    strategy_name: str
    component: str  # human-readable name, e.g. "layer_3", "head_2_5"
    description: str
    metadata: dict[str, Any] | None = None


class AblationStrategy(abc.ABC):
    """Base class that all ablation strategies must implement."""

    name: str = "base"

    @abc.abstractmethod
    def enumerate(self, handle: ModelHandle, **kwargs) -> list[AblationSpec]:
        """Return every possible ablation this strategy can perform on the model."""

    @abc.abstractmethod
    def apply(self, handle: ModelHandle, spec: AblationSpec) -> None:
        """Apply a single ablation in-place.  The caller is responsible for
        calling handle.restore() afterwards to undo the modification."""

    def iterate(self, handle: ModelHandle, **kwargs) -> Iterator[AblationSpec]:
        """Convenience: yield specs one at a time, applying + restoring around each."""
        specs = self.enumerate(handle, **kwargs)
        for spec in specs:
            self.apply(handle, spec)
            yield spec
            handle.restore()
