state_manager.py

"""
state_manager.py
----------------
Slot-filling state machine for PayParse.

Given the persistent `DialogueState` and a fresh `NLUResult`, this module:
1. Intercepts ya/tidak answers when a transaction awaits confirmation.
2. Merges newly extracted entities (non-destructively).
3. Handles mid-dialogue intent switches.
4. Resolves `target_kontak` (contact name / pronoun) into a real phone number
   BEFORE deciding which slots are still missing.
5. Computes which required slots are still missing (`REQUIRED_SLOTS`).
6. Produces the assistant reply: a follow-up question (template-based, with
   optional LLM paraphrase — "hybrid" mode), a confirmation summary, or the
   final transaction payload once the user says "ya".
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass
from typing import Optional

from schema import (
    REQUIRED_SLOTS,
    DialogueState,
    DialogueStatus,
    IntentType,
    NLUResult,
    TransactionEntities,
)

logger = logging.getLogger(__name__)

MAX_TURNS = 8  # guard against endless clarification loops

# ---------------------------------------------------------------------------
# Contact directory (stub for a real address-book / contacts service)
# Keys MUST be lowercase; lookups are done on `target_kontak.lower()`.
# ---------------------------------------------------------------------------
DUMMY_CONTACTS: dict[str, str] = {
    "anton": "081222333444",
    "budi": "08555666777",
    "nomorku": "08999888777",
    "nomor ini": "08999888777",
    "nomor saya": "08999888777",
}

# Stub GoPay registration check — in production this would query the GoPay API.
# Numbers in DUMMY_CONTACTS are "registered GoPay users"; others are not.
_GOPAY_REGISTERED_NUMBERS: set[str] = set(DUMMY_CONTACTS.values())


def check_gopay_status(phone_number: str) -> str:
    """Return 'Terdaftar GoPay' or 'Belum terdaftar GoPay' for a number."""
    if phone_number in _GOPAY_REGISTERED_NUMBERS:
        return "Terdaftar GoPay"
    return "Belum terdaftar GoPay"

# Affirmative / negative tokens for the ya-tidak confirmation intercept
_AFFIRMATIVE = {
    "ya", "iya", "y", "yes", "yoi", "yup", "oke", "ok", "okay", "sip", "siap",
    "lanjut", "lanjutkan", "gas", "boleh", "benar", "betul", "setuju", "bener",
}
_NEGATIVE = {
    "tidak", "ga", "gak", "engga", "enggak", "nggak", "ngga", "no", "n",
    "batal", "batalkan", "cancel", "jangan", "stop", "salah", "bukan",
}
# Filler / politeness words ignored when judging an affirmative answer, so
# "ya udah" or "boleh dong" still count as a clear yes.
_FILLER = {
    "hmm", "hm", "eh", "em", "ehm", "oh", "anu", "nya", "itu",
    "dong", "sih", "deh", "nih", "aja", "saja", "udah", "sudah", "langsung",
    "makasih", "terima", "kasih", "kak", "bro", "gan", "min", "pak", "bu",
}

# ---------------------------------------------------------------------------
# Follow-up question templates per (intent, slot) — natural Indonesian
# ---------------------------------------------------------------------------
FOLLOW_UP_TEMPLATES: dict[tuple[IntentType, str], str] = {
    (IntentType.TRANSFER_UANG, "recipient"): "Mau transfer ke siapa?",
    (IntentType.TRANSFER_UANG, "recipient_phone"): "Berapa nomor GoPay {recipient}?",
    (IntentType.TRANSFER_UANG, "amount"): "Berapa jumlah yang mau ditransfer?",
    (IntentType.BELI_PULSA, "phone_number"): "Pulsanya mau diisi ke nomor HP berapa?",
    (IntentType.BELI_PULSA, "amount"): "Mau beli pulsa nominal berapa?",
    (IntentType.BAYAR_PLN, "customer_id"): "Boleh sebutkan nomor ID pelanggan PLN-nya?",
    (IntentType.PESAN_GOJEK, "tujuan"): "Mau pesan Gojek ke mana tujuannya?",
    (IntentType.PESAN_GOFOOD, "makanan"): "Mau pesan makanan apa di GoFood?",
}

CONTACT_NOT_FOUND_TEMPLATE = (
    "Maaf, nomor untuk {kontak} tidak ditemukan. Berapa nomor angkanya?"
)
CONTACT_NOT_FOUND_GOPAY_TEMPLATE = (
    "Maaf, kontak {kontak} tidak ditemukan di daftar kontakmu. "
    "Berapa nomor GoPay {kontak}?"
)

CONFIRM_PROMPT_RETRY = (
    "Maaf, aku belum paham. Balas 'ya' untuk melanjutkan atau 'tidak' untuk "
    "membatalkan transaksi ini."
)

CANCEL_REPLY = "Oke, transaksinya dibatalkan. Ada lagi yang bisa aku bantu?"

UNKNOWN_REPLY = (
    "Maaf, aku belum bisa bantu untuk permintaan itu. Saat ini aku bisa bantu "
    "transfer uang, beli pulsa, bayar tagihan listrik PLN, pesan Gojek, atau "
    "pesan GoFood. Ada yang mau dibantu dari salah satunya?"
)

MAX_TURNS_REPLY = (
    "Sepertinya kita belum ketemu detail transaksinya. Yuk mulai lagi dari awal — "
    "coba sebutkan lengkap, misalnya: 'transfer 50 ribu ke Budi'."
)

INTENT_LABELS = {
    IntentType.TRANSFER_UANG: "Transfer Uang",
    IntentType.BELI_PULSA: "Beli Pulsa",
    IntentType.BAYAR_PLN: "Bayar Listrik PLN",
    IntentType.PESAN_GOJEK: "Gojek",
    IntentType.PESAN_GOFOOD: "GoFood",
}


def format_rupiah(amount: int) -> str:
    """1500000 -> 'Rp1.500.000'"""
    return f"Rp{amount:,}".replace(",", ".")


@dataclass
class TurnResult:
    """Outcome of processing one user turn."""

    state: DialogueState
    reply: str
    is_complete: bool
    payload: Optional[dict] = None
    needs_contact_resolution: bool = False
    target_kontak: Optional[str] = None


@dataclass
class _ContactResolution:
    """Internal result of _resolve_contact."""

    reply: str
    needs_client_resolution: bool = False
    contact_name: Optional[str] = None


class StateManager:
    """Evaluates dialogue state and drives the slot-filling loop.

    Parameters
    ----------
    nlu_engine:
        Only needed when `use_llm_paraphrase=True`; used to rephrase
        templated follow-up questions into more natural language.
    use_llm_paraphrase:
        Hybrid mode switch. False (default) = deterministic templates,
        zero extra API calls. True = each follow-up is paraphrased by Gemini.
    """

    def __init__(self, nlu_engine=None, use_llm_paraphrase: bool = False):
        self._nlu_engine = nlu_engine
        self._use_llm_paraphrase = use_llm_paraphrase and nlu_engine is not None

    # ------------------------------------------------------------------
    # Main entry point
    # ------------------------------------------------------------------
    def process(
        self,
        state: DialogueState,
        nlu_result: NLUResult,
        user_text: str,
        resolve_contacts_locally: bool = False,
    ) -> TurnResult:
        state.turn_count += 1
        state.history.append(user_text)

        # --- Ya/Tidak confirmation intercept (highest priority) -------
        # When a transaction is awaiting confirmation, this turn is a decision,
        # not a new extraction — handle it before any slot logic runs.
        if state.status == DialogueStatus.CONFIRMING:
            decision = self._classify_confirmation(user_text)
            if decision is True:
                return self._finalize(state)
            if decision is False:
                self._reset(state)
                return TurnResult(state, CANCEL_REPLY, is_complete=False)
            # Neither ya nor tidak: allow a genuine correction/new request to
            # fall through, otherwise re-ask for a clear decision.
            has_correction = (
                nlu_result.intent != IntentType.UNKNOWN
                and bool(nlu_result.entities.filled_fields())
            )
            if not has_correction:
                return TurnResult(state, CONFIRM_PROMPT_RETRY, is_complete=False)

        # --- Out-of-scope handling -----------------------------------
        if nlu_result.intent == IntentType.UNKNOWN:
            if state.intent == IntentType.UNKNOWN:
                # No active transaction: politely state our scope.
                state.status = DialogueStatus.UNKNOWN
                return TurnResult(state, UNKNOWN_REPLY, is_complete=False)
            # Active transaction, but this turn was unintelligible:
            # re-ask for the pending slot instead of derailing.
            return self._ask_next_slot(state)

        # --- Intent switch mid-dialogue ------------------------------
        if (
            state.intent != IntentType.UNKNOWN
            and nlu_result.intent != state.intent
        ):
            logger.info(
                "Intent switch: %s -> %s (resetting entities)",
                state.intent.value,
                nlu_result.intent.value,
            )
            state.entities = TransactionEntities()

        state.intent = nlu_result.intent

        # --- Detect contact name change → clear old resolved number ---
        # If the user mentions a NEW target_kontak (different name from the
        # current one), the previously resolved phone number is stale and
        # must be cleared so the new contact gets its own resolution.
        new_kontak = nlu_result.entities.target_kontak
        old_kontak = state.entities.target_kontak
        if (
            new_kontak is not None
            and old_kontak is not None
            and new_kontak.lower() != old_kontak.lower()
        ):
            logger.info(
                "Contact changed: '%s' -> '%s' (clearing old phone number)",
                old_kontak, new_kontak,
            )
            state.entities.recipient_phone = None
            state.entities.phone_number = None

        state.entities = state.entities.merge(nlu_result.entities)

        # --- Contact resolution ---------------------------------------
        # Runs BEFORE the missing-slot check so a resolvable contact name
        # ("beliin anton pulsa") never triggers a redundant "nomor berapa?".
        unresolved = self._resolve_contact(state, resolve_contacts_locally)
        if unresolved is not None:
            if unresolved.needs_client_resolution:
                state.status = DialogueStatus.COLLECTING
                state.missing_slots = self._compute_missing(state)
                return TurnResult(
                    state,
                    unresolved.reply,
                    is_complete=False,
                    needs_contact_resolution=True,
                    target_kontak=unresolved.contact_name,
                )
            state.status = DialogueStatus.COLLECTING
            state.missing_slots = self._compute_missing(state)
            return TurnResult(state, unresolved.reply, is_complete=False)

        state.missing_slots = self._compute_missing(state)

        # --- Loop guard ----------------------------------------------
        if state.missing_slots and state.turn_count >= MAX_TURNS:
            self._reset(state)
            return TurnResult(state, MAX_TURNS_REPLY, is_complete=False)

        # --- Ask for confirmation or keep collecting ------------------
        if not state.missing_slots:
            return self._confirm(state)
        return self._ask_next_slot(state)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _compute_missing(state: DialogueState) -> list[str]:
        required = REQUIRED_SLOTS[state.intent]
        filled = state.entities.filled_fields()
        return [slot for slot in required if slot not in filled]

    def _ask_next_slot(self, state: DialogueState) -> TurnResult:
        state.status = DialogueStatus.COLLECTING
        slot = state.missing_slots[0]
        question = FOLLOW_UP_TEMPLATES.get(
            (state.intent, slot), f"Boleh lengkapi informasi '{slot}'?"
        )
        # Fill in template placeholders (e.g. {recipient} for GoPay number ask)
        if state.entities.recipient and "{recipient}" in question:
            question = question.format(recipient=state.entities.recipient)
        if self._use_llm_paraphrase:
            question = self._nlu_engine.paraphrase(question)
        return TurnResult(state, question, is_complete=False)

    @staticmethod
    def _resolve_contact(
        state: DialogueState, client_resolves: bool = False
    ) -> Optional["_ContactResolution"]:
        """Resolve `target_kontak` (a name or pronoun) into a real number.

        Handles both beli_pulsa (→ phone_number) and transfer_uang (→ recipient_phone).
        Returns None when nothing needs resolving or resolution succeeded.
        Otherwise returns a _ContactResolution describing the follow-up:
        - When client_resolves is True, signals the client to resolve
          on-device (needs_client_resolution=True).
        - When False (web/demo), tries DUMMY_CONTACTS first, then asks
          for raw digits if the contact is unknown.
        """
        entities = state.entities
        intent = state.intent

        # Determine which slot needs a phone number and whether it's already filled
        if intent == IntentType.BELI_PULSA:
            phone_slot = "phone_number"
            current_phone = entities.phone_number
            not_found_template = CONTACT_NOT_FOUND_TEMPLATE
        elif intent == IntentType.TRANSFER_UANG:
            phone_slot = "recipient_phone"
            current_phone = entities.recipient_phone
            not_found_template = CONTACT_NOT_FOUND_GOPAY_TEMPLATE
        else:
            return None

        if entities.target_kontak is None or current_phone is not None:
            return None

        # Mobile flow: delegate to the client immediately.
        if client_resolves:
            logger.info(
                "Asking client to resolve contact '%s'", entities.target_kontak
            )
            return _ContactResolution(
                reply=(
                    f"Boleh pilih kontak {entities.target_kontak} dari daftar "
                    f"kontakmu?"
                ),
                needs_client_resolution=True,
                contact_name=entities.target_kontak,
            )

        # Web/demo flow: try the built-in directory first.
        key = entities.target_kontak.lower()
        number = DUMMY_CONTACTS.get(key)
        if number:
            setattr(entities, phone_slot, number)
            logger.info("Resolved contact '%s' -> %s (slot=%s)", key, number, phone_slot)
            return None

        logger.info("Contact '%s' not found in directory", key)
        return _ContactResolution(
            reply=not_found_template.format(kontak=entities.target_kontak),
            needs_client_resolution=False,
        )

    @staticmethod
    def _classify_confirmation(user_text: str) -> Optional[bool]:
        """True = confirmed, False = cancelled, None = unclear.

        Asymmetric on purpose (fail-safe for money movement):
        - Any negative token cancels — over-cancelling is harmless.
        - An affirmative only counts if the WHOLE utterance is affirmative
          (after dropping fillers). Otherwise Indonesian fillers like
          "apa ya" / "gimana ya" would falsely confirm a transaction.
        """
        tokens = set(re.findall(r"[a-z]+", user_text.lower()))
        if tokens & _NEGATIVE:
            return False
        meaningful = tokens - _FILLER
        if meaningful and meaningful <= _AFFIRMATIVE:
            return True
        return None

    def _confirm(self, state: DialogueState) -> TurnResult:
        """All slots filled: show the summary and WAIT for ya/tidak.
        The payload is withheld until the user actually confirms."""
        state.status = DialogueStatus.CONFIRMING
        return TurnResult(
            state, self._build_confirmation_text(state), is_complete=False
        )

    def _finalize(self, state: DialogueState) -> TurnResult:
        """User said ya: release the transaction payload."""
        state.status = DialogueStatus.COMPLETE
        payload = self._build_payload(state)
        label = INTENT_LABELS[state.intent]

        # Gojek/GoFood: show "mencari driver" + "menunggu driver" message
        if state.intent == IntentType.PESAN_GOJEK:
            reply = (
                f"Mencari driver terdekat untuk rute "
                f"{state.entities.asal or 'Bogor'} - {state.entities.tujuan}...\n"
                f"Silakan menunggu driver Budi (default). "
                f"Transaksi {label} berhasil diproses. Terima kasih!"
            )
        elif state.intent == IntentType.PESAN_GOFOOD:
            reply = (
                f"Mencari driver terdekat untuk pesanan {state.entities.makanan}...\n"
                f"Silakan menunggu driver Budi (default). "
                f"Transaksi {label} berhasil diproses. Terima kasih!"
            )
        else:
            reply = f"Transaksi {label} berhasil diproses. Terima kasih!"

        return TurnResult(
            state,
            reply,
            is_complete=True,
            payload=payload,
        )

    @staticmethod
    def _build_payload(state: DialogueState) -> dict:
        """Machine-consumable transaction payload for downstream services."""
        entities = state.entities.model_dump(exclude_none=True)
        return {
            "transaction_type": state.intent.value,
            "entities": entities,
            "turns_to_complete": state.turn_count,
        }

    @staticmethod
    def _build_confirmation_text(state: DialogueState) -> str:
        e = state.entities
        label = INTENT_LABELS[state.intent]
        lines = [f"Siap! Konfirmasi {label}:"]
        if state.intent == IntentType.TRANSFER_UANG:
            nama = e.recipient or "(tanpa nama)"
            nomor = e.recipient_phone or "?"
            gopay = check_gopay_status(nomor)
            lines.append(f"- Nama     : {nama}")
            lines.append(f"- Nomor    : {nomor}")
            lines.append(f"- Status   : {gopay}")
            lines.append(f"- Jumlah   : {format_rupiah(e.amount)}")
        elif state.intent == IntentType.BELI_PULSA:
            nomor = e.phone_number
            if e.target_kontak:
                nomor = f"{e.phone_number} ({e.target_kontak})"
            lines.append(f"- Nomor HP : {nomor}")
            lines.append(f"- Nominal  : {format_rupiah(e.amount)}")
            if e.provider:
                lines.append(f"- Provider : {e.provider}")
        elif state.intent == IntentType.BAYAR_PLN:
            lines.append(f"- ID Pelanggan : {e.customer_id}")
            if e.amount:
                lines.append(f"- Nominal      : {format_rupiah(e.amount)}")
        elif state.intent == IntentType.PESAN_GOJEK:
            asal = e.asal or "Bogor"
            lines.append(f"- Asal    : {asal}")
            lines.append(f"- Tujuan  : {e.tujuan}")
        elif state.intent == IntentType.PESAN_GOFOOD:
            lines.append(f"- Pesanan : {e.makanan}")
        lines.append("Lanjutkan?")
        return "\n".join(lines)

    @staticmethod
    def _reset(state: DialogueState) -> None:
        state.intent = IntentType.UNKNOWN
        state.entities = TransactionEntities()
        state.missing_slots = []
        state.status = DialogueStatus.UNKNOWN
        state.turn_count = 0