from __future__ import annotations

from enum import Enum
from typing import Callable, Optional, TypeVar

from domain.data_transformation.column_classifier import ColumnClassifier

E = TypeVar("E", bound=Enum)


class GuardedColumnClassifier(ColumnClassifier[E]):
    """A ``ColumnClassifier`` that resolves the descriptions a deterministic guard
    is certain about, and delegates the rest to a fallback classifier.

    The ``guard`` maps a raw description to a category member when it recognises it
    deterministically (e.g. a party-ceiling roof marker — #1376), else ``None``.
    Guard hits never reach the fallback, so an unreliable classifier (the LLM)
    cannot override a description the guard is sure of — and the LLM is not billed
    for it. Every description still appears in the result (guarded or fallen-back).
    """

    def __init__(
        self,
        guard: Callable[[str], Optional[E]],
        fallback: ColumnClassifier[E],
    ) -> None:
        self._guard = guard
        self._fallback = fallback

    def classify(self, descriptions: set[str]) -> dict[str, E]:
        guarded: dict[str, E] = {}
        misses: set[str] = set()
        for description in descriptions:
            member = self._guard(description)
            if member is not None:
                guarded[description] = member
            else:
                misses.add(description)
        # Only the misses reach the fallback — a fully-guarded batch never calls it.
        if misses:
            guarded.update(self._fallback.classify(misses))
        return guarded