from __future__ import annotations

from typing import Iterable, Mapping


DEFAULT_STAGE_POLICY = {
    "ARGUMENTATION": True,
    "DATA_CONFIRMATION": False,
    "FORMALIZATION": False,
    "PRESENTATION": True,
    "DONE": False,
    "INTRO": True,
    "IDLE_NUDGE": False,
    "AGENT_WAIT_TIMEOUT_RETRY": True,
}

DEFAULT_BACKCHANNEL_TOKENS = {
    "uhum",
    "aham",
    "hm",
    "hmm",
    "ok",
    "certo",
    "tá",
    "ta",
    "sim",
}


class InterruptPolicy:
    def __init__(
        self,
        *,
        stage_policy: Mapping[str, bool] | None = None,
        backchannel_tokens: Iterable[str] | None = None,
    ) -> None:
        self._stage_policy = {
            str(key).strip().upper(): bool(value)
            for key, value in (stage_policy or DEFAULT_STAGE_POLICY).items()
        }
        self._backchannel_tokens = {
            str(token).strip().lower()
            for token in (backchannel_tokens or DEFAULT_BACKCHANNEL_TOKENS)
            if str(token).strip()
        }

    def allow_stage(self, stage: str) -> bool:
        key = (stage or "").strip().upper()
        return bool(self._stage_policy.get(key, True))

    def should_ignore_backchannel(self, *, speaking: bool, text: str) -> bool:
        if not speaking:
            return False
        if len(text) <= 2:
            return True
        lowered = (text or "").strip().lower()
        return lowered in self._backchannel_tokens
