from collections.abc import Callable
from dataclasses import dataclass, field

from .filters import strip_fillers
from .rules import MARKERS

LEVELS = {
    "light": ["zero_copula"],
    "medium": ["zero_copula", "prospective_finna", "continuative_steady"],
    "full": [
        "zero_copula",
        "habitual_be",
        "completive_done",
        "prospective_finna",
        "continuative_steady",
    ],
}

_APPLY_ORDER = (
    "habitual_be",
    "completive_done",
    "prospective_finna",
    "continuative_steady",
    "zero_copula",
)


@dataclass
class CompressionResult:
    original: str
    compressed: str
    level: str
    markers_applied: list[str] = field(default_factory=list)
    filler_stripped: bool = False
    gated: bool = False


def compress(
    text: str,
    level: str = "medium",
    *,
    strip_filler: bool = False,
    tokenizer: Callable[[str], int] | None = None,
    gate: bool = False,
) -> CompressionResult:
    """Apply the marker transforms for `level` in sequence and return a CompressionResult.

    Keyword-only options (defaults preserve the original behavior exactly):

    * ``strip_filler=True`` runs :func:`aave_compress.filters.strip_fillers`
      as the FIRST pass, before any AAVE marker — the generic
      low-information-word mechanism layer under the AAVE grammar layer.
    * ``tokenizer`` — optional callable ``str -> int`` returning a token
      count, e.g. ``lambda s: count_tokens_tiktoken(s, "gpt-4o").tokens``.
    * ``gate=True`` (requires ``tokenizer``) makes every transform step —
      the filler strip and each individual AAVE marker — independently
      gated: a step is kept only if it STRICTLY reduces the token count
      under ``tokenizer``; otherwise it is reverted. This guarantees the
      compressor never makes a prompt bigger (as measured by ``tokenizer``).

    ``markers_applied`` (and ``filler_stripped``) reflect only steps that
    actually changed the text AND survived the gate.
    """
    if level not in LEVELS:
        raise ValueError(f"unknown level {level!r}; expected one of {sorted(LEVELS)}")
    if gate and tokenizer is None:
        raise ValueError(
            "gate=True requires a tokenizer: pass tokenizer=<callable str -> int> "
            "(e.g. lambda s: count_tokens_tiktoken(s, 'gpt-4o').tokens). "
            "Gating cannot measure token effects without a way to count."
        )
    enabled = set(LEVELS[level])
    compressed = text
    current_tokens = tokenizer(text) if gate else None
    markers_applied: list[str] = []
    filler_stripped = False

    def apply_step(
        current: str, transform: Callable[[str], str]
    ) -> tuple[str, bool]:
        nonlocal current_tokens
        candidate = transform(current)
        if candidate == current:
            return current, False
        if gate:
            candidate_tokens = tokenizer(candidate)
            if candidate_tokens >= current_tokens:
                return current, False  # revert: no strict token win
            current_tokens = candidate_tokens
        return candidate, True

    if strip_filler:
        compressed, filler_stripped = apply_step(compressed, strip_fillers)

    for name in _APPLY_ORDER:
        if name not in enabled:
            continue
        compressed, changed = apply_step(compressed, MARKERS[name])
        if changed:
            markers_applied.append(name)

    return CompressionResult(
        original=text,
        compressed=compressed,
        level=level,
        markers_applied=markers_applied,
        filler_stripped=filler_stripped,
        gated=gate,
    )