nlu_engine.py

"""
nlu_engine.py
-------------
LLM-based NLU engine for PayParse.

- `NLUEngine`     : zero-shot Gemini Flash with strict Pydantic structured
                    output (`response_schema=NLUResult`), slang-aware prompt,
                    dialogue-context injection, retries, graceful fallback.
- `MockNLUEngine` : deterministic rule-based engine with the same interface,
                    used for offline evaluation / running without an API key.
"""

from __future__ import annotations

import logging
import os
import re
import time
from typing import Optional, Protocol

from dotenv import load_dotenv

from schema import DialogueState, IntentType, NLUResult, TransactionEntities
from ner_extractor import ner_extractor, NERResult

load_dotenv()
logger = logging.getLogger(__name__)

MODEL_NAME = os.getenv("PAYPARSE_MODEL", "gemini-flash-lite-latest")


# ---------------------------------------------------------------------------
# NER + LLM merge: NER fills gaps the LLM missed, LLM wins on conflicts.
# ---------------------------------------------------------------------------
def _merge_ner_with_llm(ner: NERResult, llm: NLUResult) -> NLUResult:
    """Merge NER pre-scan with LLM output.

    Strategy:
    - Intent: LLM wins (more context-aware). NER only used as fallback.
    - Amount: NER wins if LLM missed it (NER is deterministic for amounts).
    - Phone/recipient_phone: NER wins if LLM missed it.
    - recipient/target_kontak: LLM wins (name extraction is ambiguous).
    - customer_id: NER wins if LLM missed it.
    - provider: NER wins if LLM missed it.
    """
    e_llm = llm.entities
    fields = ner.extracted_fields

    # Amount: NER fills gap
    if "amount" in fields and e_llm.amount is None and ner.amount is not None:
        e_llm.amount = ner.amount
        logger.info("NER filled amount=%d", ner.amount)

    # Phone number: NER fills gap
    if "phone_number" in fields and e_llm.phone_number is None and ner.phone_number is not None:
        e_llm.phone_number = ner.phone_number
        logger.info("NER filled phone_number=%s", ner.phone_number)

    # Recipient phone: NER fills gap
    if "recipient_phone" in fields and e_llm.recipient_phone is None and ner.recipient_phone is not None:
        e_llm.recipient_phone = ner.recipient_phone
        logger.info("NER filled recipient_phone=%s", ner.recipient_phone)

    # Customer ID: NER fills gap
    if "customer_id" in fields and e_llm.customer_id is None and ner.customer_id is not None:
        e_llm.customer_id = ner.customer_id
        logger.info("NER filled customer_id=%s", ner.customer_id)

    # Provider: NER fills gap
    if "provider" in fields and e_llm.provider is None and ner.provider is not None:
        e_llm.provider = ner.provider
        logger.info("NER filled provider=%s", ner.provider)

    # Target kontak: NER fills gap (only if LLM didn't extract it and no phone)
    if (
        "target_kontak" in fields
        and e_llm.target_kontak is None
        and ner.target_kontak is not None
        and e_llm.phone_number is None
        and e_llm.recipient_phone is None
    ):
        e_llm.target_kontak = ner.target_kontak
        if e_llm.recipient is None and ner.recipient is not None:
            e_llm.recipient = ner.recipient
        logger.info("NER filled target_kontak=%s", ner.target_kontak)

    # Recipient: NER fills gap
    if "recipient" in fields and e_llm.recipient is None and ner.recipient is not None:
        e_llm.recipient = ner.recipient
        logger.info("NER filled recipient=%s", ner.recipient)

    # Tujuan (Gojek destination): NER fills gap
    if "tujuan" in fields and e_llm.tujuan is None and ner.tujuan is not None:
        e_llm.tujuan = ner.tujuan
        logger.info("NER filled tujuan=%s", ner.tujuan)

    # Asal (Gojek origin): NER fills gap
    if "asal" in fields and e_llm.asal is None and ner.asal is not None:
        e_llm.asal = ner.asal
        logger.info("NER filled asal=%s", ner.asal)

    # Makanan (GoFood item): NER fills gap
    if "makanan" in fields and e_llm.makanan is None and ner.makanan is not None:
        e_llm.makanan = ner.makanan
        logger.info("NER filled makanan=%s", ner.makanan)

    # Bump confidence if NER agrees with LLM
    if ner.intent is not None and ner.intent == llm.intent:
        llm.confidence = min(1.0, llm.confidence + 0.05)

    return llm

# ---------------------------------------------------------------------------
# System prompt: the heart of the zero-shot NLU
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = """\
Kamu adalah mesin NLU (Natural Language Understanding) untuk asisten suara
dompet digital Indonesia (seperti GoPay). Tugasmu HANYA mengekstrak intent
dan entities dari ucapan pengguna, lalu mengembalikan JSON sesuai skema.

## INTENT yang didukung:
1. "transfer_uang" : kirim/transfer uang ke seseorang.
   Kata kunci: transfer, tf, kirim, kirimin, transferin, kasih uang.
2. "beli_pulsa"    : beli pulsa / isi ulang / top-up pulsa ke nomor HP.
   Kata kunci: pulsa, isi pulsa, top up, isiin.
3. "bayar_pln"     : bayar tagihan listrik PLN / beli token listrik.
   Kata kunci: listrik, PLN, token, tagihan listrik, tagihan pln, meteran.
4. "pesan_gojek"   : pesan ojek/ride Gojek dari asal ke tujuan.
   Kata kunci: gojek, go jek, goride, naik gojek, pesan gojek.
   Entity "asal" default "Bogor" jika tidak disebut.
   Entity "tujuan" adalah lokasi tujuan (stasiun, bandara, mall, dll).
5. "pesan_gofood"  : pesan makanan via GoFood.
   Kata kunci: gofood, go food, pesan makan, beli makan.
   Entity "makanan" adalah nama makanan (nasi goreng, ayam geprek, dll).
6. "unknown"       : semua permintaan di luar 5 intent di atas
   (contoh: tanya cuaca, ngobrol basa-basi).

## NORMALISASI SLANG UANG (WAJIB dikonversi ke integer Rupiah):
- "seceng" / "seribu"      = 1000
- "goceng"                 = 5000
- "ceban"                  = 10000
- "noban"                  = 20000
- "gocap" / "gopek ribu"?  -> "gocap" = 50000, "gopek" = 500
- "cepek"                  = 100 (uang: biasanya maksudnya "cepek ribu" = 100000
  jika konteksnya transfer/pulsa; gunakan 100000 untuk konteks transaksi)
- "seket" (Jawa)           = 50 -> "seket ewu" = 50000
- "sejuta" / "1jt" / "1 juta" = 1000000
- "50rb" / "50ribu" / "50k"   = 50000
- "2,5jt" / "2.5 juta"        = 2500000

## ATURAN PENTING:
- Pengguna sering typo ("pusla" = pulsa, "trasnfer" = transfer) dan memakai
  bahasa lisan tidak baku. Tetap pahami maksudnya.
- HANYA ekstrak entity yang DISEBUT EKSPLISIT oleh pengguna.
  JANGAN PERNAH mengarang/menebak nilai yang tidak diucapkan.
  Jika tidak disebut, biarkan null.
- "recipient" adalah nama orang penerima transfer (budi, mama, bang jono).
- "recipient_phone" adalah nomor HP penerima transfer, HANYA jika pengguna
  menyebutkan DIGIT ANGKA secara eksplisit. Jangan pernah mengarang digit.
- "phone_number" adalah nomor HP tujuan pulsa, HANYA jika pengguna menyebutkan
  DIGIT ANGKA secara eksplisit.
- "target_kontak" dipakai jika pengguna merujuk nomor HP secara TIDAK LANGSUNG:
  * lewat nama kontak: "beliin anton pulsa" -> target_kontak = "anton"
  * lewat nama kontak transfer: "tf ke budi" -> recipient = "budi", target_kontak = "budi"
  * lewat kata ganti  : "isi ke nomor ini", "nomorku" -> target_kontak = "nomor ini" / "nomorku"
  Dalam kasus ini "phone_number" dan "recipient_phone" WAJIB null — jangan
  pernah mengarang digit. Sebaliknya, jika digit sudah disebut,
  "target_kontak" biarkan null.

  Contoh:
  Ucapan: "beliin anton pulsa 10rb"
  -> {"intent": "beli_pulsa", "entities": {"amount": 10000,
      "phone_number": null, "target_kontak": "anton"}}
  Ucapan: "tf 50rb ke budi"
  -> {"intent": "transfer_uang", "entities": {"recipient": "budi",
      "recipient_phone": null, "target_kontak": "budi", "amount": 50000}}
  Ucapan: "transfer ke 081234567890 100rb"
  -> {"intent": "transfer_uang", "entities": {"recipient_phone": "081234567890",
      "recipient": null, "target_kontak": null, "amount": 100000}}
  Ucapan: "gojek ke stasiun"
  -> {"intent": "pesan_gojek", "entities": {"asal": null, "tujuan": "stasiun"}}
  Ucapan: "gofood nasi goreng"
  -> {"intent": "pesan_gofood", "entities": {"makanan": "nasi goreng"}}
- "customer_id" adalah nomor ID pelanggan / meteran PLN (hanya digit).
- Nomor yang diawali 08 kemungkinan besar phone_number, bukan customer_id.
- Jika pengguna sedang menjawab pertanyaan lanjutan (lihat KONTEKS DIALOG),
  jawaban singkat seperti "buat budi" atau "50rb" adalah pengisian slot untuk
  intent yang SEDANG BERJALAN — pertahankan intent tersebut.
- PENTING: entity yang SUDAH terisi di KONTEKS DIALOG JANGAN diulang lagi di
  output "entities" kecuali pengguna benar-benar menyebutkannya ulang di
  UCAPAN saat ini. Field "entities" HANYA berisi hal baru yang disebut di
  UCAPAN PENGGUNA sekarang, bukan salinan dari konteks.
- "normalized_text": tulis ulang ucapan dalam bahasa Indonesia baku dan rapi.
- "confidence": estimasi keyakinanmu terhadap intent (0.0 - 1.0).
"""


def _build_context_block(state: Optional[DialogueState]) -> str:
    """Render the current dialogue state so short follow-up answers
    ("buat budi", "50rb") are resolved against the ongoing intent."""
    if state is None or state.intent == IntentType.UNKNOWN:
        return ""
    filled = {
        k: v for k, v in state.entities.model_dump().items() if v is not None
    }
    return (
        "\n## KONTEKS DIALOG (percakapan sedang berjalan):\n"
        f"- Intent aktif: {state.intent.value}\n"
        f"- Entity yang sudah terisi: {filled or 'belum ada'}\n"
        f"- Slot yang masih ditanyakan: {state.missing_slots}\n"
        "Ucapan berikut kemungkinan adalah jawaban untuk slot yang ditanyakan.\n"
    )


class BaseNLUEngine(Protocol):
    """Common interface so the API/eval can swap live and mock engines."""

    def extract(self, text: str, state: Optional[DialogueState] = None) -> NLUResult:
        ...


# ---------------------------------------------------------------------------
# Live Gemini engine
# ---------------------------------------------------------------------------
class NLUEngine:
    """Gemini-backed NLU with strict structured output.

    Supports two SDKs:
    - google-genai (new SDK, preferred — has response_schema for Pydantic)
    - google-generativeai (old SDK, fallback — needed for HF Spaces where
      google-genai's websockets>=13 conflicts with gradio-client's websockets<13)
    """

    def __init__(self, api_key: Optional[str] = None, max_retries: int = 2):
        api_key = api_key or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
        if not api_key:
            raise ValueError(
                "GEMINI_API_KEY not set. Export it or add it to a .env file."
            )
        self._max_retries = max_retries
        self._sdk = None
        self._client = None
        self._model = None

        # Try new SDK (google-genai) first
        try:
            from google import genai
            self._client = genai.Client(api_key=api_key)
            self._sdk = "genai"
            logger.info("Using google-genai SDK")
        except ImportError:
            pass

        # Fall back to old SDK (google-generativeai)
        if self._sdk is None:
            try:
                import google.generativeai as genai
                # Old SDK needs GOOGLE_API_KEY env var or configure()
                os.environ["GOOGLE_API_KEY"] = api_key
                genai.configure(api_key=api_key)
                self._model = genai.GenerativeModel(
                    MODEL_NAME, system_instruction=SYSTEM_PROMPT
                )
                self._sdk = "generativeai"
                logger.info("Using google-generativeai SDK")
            except ImportError:
                raise ImportError(
                    "Neither google-genai nor google-generativeai is installed. "
                    "Install one: pip install google-genai OR pip install google-generativeai"
                )

    def extract(self, text: str, state: Optional[DialogueState] = None) -> NLUResult:
        """Run zero-shot extraction with NER pre-scan.

        Pipeline: NER pre-scan → LLM extraction → merge (NER fills gaps
        the LLM missed, LLM wins on conflicts for ambiguous fields).

        Returns an `unknown` NLUResult on failure so the pipeline degrades
        gracefully instead of crashing.
        """
        # --- NER pre-scan (rule-based, zero latency) ---
        ner = ner_extractor.extract(text)
        logger.info("NER pre-scan: %s", ner.extracted_fields)

        prompt = _build_context_block(state) + f'\n## UCAPAN PENGGUNA:\n"{text}"'

        for attempt in range(self._max_retries + 1):
            try:
                llm_result = self._call_llm(prompt)
                return _merge_ner_with_llm(ner, llm_result)
            except Exception as exc:  # network, rate limit, malformed JSON
                logger.warning("NLU attempt %d failed: %s", attempt + 1, exc)
                if attempt < self._max_retries:
                    time.sleep(2 ** attempt)  # 1s, 2s backoff

        # --- Fallback: use NER result if LLM fails entirely ---
        if ner.intent is not None:
            logger.info("LLM failed — falling back to NER result.")
            return NLUResult(
                intent=ner.intent,
                entities=ner.to_entities(),
                confidence=ner.confidence,
                normalized_text=text,
            )

        logger.error("NLU extraction failed after retries; returning unknown.")
        return NLUResult(intent=IntentType.UNKNOWN, normalized_text=text)

    def _call_llm(self, prompt: str) -> NLUResult:
        """Call the LLM using whichever SDK is available."""
        if self._sdk == "genai":
            return self._call_genai(prompt)
        else:
            return self._call_generativeai(prompt)

    def _call_genai(self, prompt: str) -> NLUResult:
        """Call via google-genai (new SDK with response_schema)."""
        from google.genai import types
        config = types.GenerateContentConfig(
            system_instruction=SYSTEM_PROMPT,
            response_mime_type="application/json",
            response_schema=NLUResult,
            temperature=0.0,
        )
        response = self._client.models.generate_content(
            model=MODEL_NAME, contents=prompt, config=config
        )
        return response.parsed or NLUResult.model_validate_json(response.text)

    def _call_generativeai(self, prompt: str) -> NLUResult:
        """Call via google-generativeai (old SDK, manual JSON parse)."""
        import google.generativeai as genai
        response = self._model.generate_content(
            prompt,
            generation_config=genai.GenerationConfig(
                temperature=0.0,
                response_mime_type="application/json",
            ),
        )
        return NLUResult.model_validate_json(response.text)

    def paraphrase(self, question: str) -> str:
        """Optionally rephrase a templated follow-up question into a more
        natural, friendly sentence (hybrid follow-up generation)."""
        try:
            if self._sdk == "genai":
                response = self._client.models.generate_content(
                    model=MODEL_NAME,
                    contents=(
                        "Tulis ulang pertanyaan asisten dompet digital berikut agar "
                        "terdengar ramah dan natural dalam bahasa Indonesia santai. "
                        "Balas HANYA dengan satu kalimat pertanyaannya saja.\n"
                        f"Pertanyaan: {question}"
                    ),
                )
                return (response.text or question).strip()
            else:
                response = self._model.generate_content(
                    "Tulis ulang pertanyaan asisten dompet digital berikut agar "
                    "terdengar ramah dan natural dalam bahasa Indonesia santai. "
                    "Balas HANYA dengan satu kalimat pertanyaannya saja.\n"
                    f"Pertanyaan: {question}"
                )
                return (response.text or question).strip()
        except Exception as exc:
            logger.warning("Paraphrase failed, using template: %s", exc)
            return question


# ---------------------------------------------------------------------------
# Offline mock engine (rule-based)
# ---------------------------------------------------------------------------
_SLANG_AMOUNTS = {
    "seceng": 1_000,
    "seribu": 1_000,
    "goceng": 5_000,
    "ceban": 10_000,
    "noban": 20_000,
    "gocap": 50_000,
    "cepek": 100_000,   # transactional context
    "sejuta": 1_000_000,
}

_INTENT_KEYWORDS = {
    IntentType.BELI_PULSA: ["pulsa", "pusla", "top up", "topup", "isi ulang", "isiin"],
    IntentType.BAYAR_PLN: ["listrik", "pln", "token", "meteran", "tagihan pln", "tagihan listrik"],
    IntentType.TRANSFER_UANG: [
        "transfer", "trasnfer", "tf", "kirim", "kirimin", "transferin", "kasih", "beri"
    ],
    IntentType.PESAN_GOJEK: ["gojek", "go jek", "goride", "go ride"],
    IntentType.PESAN_GOFOOD: ["gofood", "go food", "pesan makan", "beli makan"],
}


class MockNLUEngine:
    """Deterministic keyword/regex NLU with the same interface as `NLUEngine`.
    Good enough to exercise the state machine and run tests offline —
    NOT a substitute for the LLM's robustness."""

    def extract(self, text: str, state: Optional[DialogueState] = None) -> NLUResult:
        lowered = f" {text.lower()} "
        intent = self._classify(lowered, state)
        entities = self._extract_entities(lowered, intent)
        return NLUResult(
            intent=intent,
            entities=entities,
            confidence=0.5 if intent == IntentType.UNKNOWN else 0.9,
            normalized_text=text,
        )

    def paraphrase(self, question: str) -> str:
        return question  # mock: templates pass through unchanged

    def _classify(self, lowered: str, state: Optional[DialogueState]) -> IntentType:
        for intent, keywords in _INTENT_KEYWORDS.items():
            if any(f" {kw} " in lowered or lowered.strip().startswith(kw)
                   for kw in keywords):
                return intent
        # Short answers during slot filling keep the active intent
        if state is not None and state.intent != IntentType.UNKNOWN:
            return state.intent
        return IntentType.UNKNOWN

    def _extract_entities(
        self, lowered: str, intent: IntentType
    ) -> TransactionEntities:
        entities = TransactionEntities()
        if intent == IntentType.UNKNOWN:
            return entities

        entities.amount = self._parse_amount(lowered)

        # Phone numbers (start with 08, 9-13 digits) vs PLN customer IDs
        numbers = re.findall(r"\b(\d[\d\-\s]{7,15}\d)\b", lowered)
        for raw in numbers:
            digits = re.sub(r"\D", "", raw)
            if digits.startswith("08") and 9 <= len(digits) <= 13:
                entities.phone_number = digits
            elif intent == IntentType.BAYAR_PLN:
                entities.customer_id = digits

        if intent == IntentType.TRANSFER_UANG:
            # If raw digits were mentioned, fill recipient_phone
            if entities.phone_number is not None:
                entities.recipient_phone = entities.phone_number
                entities.phone_number = None
            # Extract recipient name
            m = re.search(
                r"\b(?:ke|buat|untuk)\s+(?:(?:si|bang|mbak|pak|bu)\s+)?([a-z]+)",
                lowered,
            )
            if m and m.group(1) not in {"nomor", "rekening", "hp"}:
                entities.recipient = m.group(1).capitalize()
                # If no raw digits, fill target_kontak for resolution
                if entities.recipient_phone is None:
                    entities.target_kontak = entities.recipient

        # Indirect phone reference: contact name or pronoun instead of digits
        if intent == IntentType.BELI_PULSA and entities.phone_number is None:
            entities.target_kontak = self._parse_contact(lowered)

        # --- Gojek: extract tujuan ---
        if intent == IntentType.PESAN_GOJEK:
            m = re.search(r"\bgojek\s+(?:ke|buat|untuk)\s+(.+?)(?:\s+dari\s|$)", lowered)
            if m:
                dest = m.group(1).strip()
                # Strip filler words
                dest = re.sub(r"\s+(?:dong|sih|deh|nih|aja|ya|yah)$", "", dest).strip()
                if dest:
                    entities.tujuan = dest.capitalize()
            else:
                m = re.search(r"\bgojek\s+([a-z][a-z\s]+)", lowered)
                if m and m.group(1).strip() not in {"ke", "dari", "pesan"}:
                    dest = m.group(1).strip()
                    dest = re.sub(r"\s+(?:dong|sih|deh|nih|aja|ya|yah)$", "", dest).strip()
                    if dest:
                        entities.tujuan = dest.capitalize()
            # Asal: "dari X" or default Bogor
            m = re.search(r"\bdari\s+([a-z][a-z\s]+?)(?:\s+ke\s|$)", lowered)
            if m:
                entities.asal = m.group(1).strip().capitalize()

        # --- GoFood: extract makanan ---
        if intent == IntentType.PESAN_GOFOOD:
            m = re.search(r"\bgofood\s+(.+?)(?:\s*$)", lowered)
            if m:
                food = m.group(1).strip()
                if food not in {"pesan", "order", "beli", "mau"}:
                    entities.makanan = food.capitalize()
            else:
                m = re.search(r"\b(?:beli|pesan)\s+makan(?:an)?\s+(.+?)(?:\s*$)", lowered)
                if m:
                    entities.makanan = m.group(1).strip().capitalize()

        return entities

    @staticmethod
    def _parse_contact(lowered: str) -> Optional[str]:
        """Detect pronouns ('nomorku', 'nomor ini') or a contact name."""
        for pronoun in ("nomor ini", "nomer ini", "nomorku", "nomerku", "nomor saya"):
            if pronoun in lowered:
                return pronoun
        # "isi ke anton", "beliin anton pulsa"
        m = re.search(r"\b(?:ke|buat|untuk)\s+([a-z]+)", lowered)
        if m and m.group(1) not in {"nomor", "nomer", "hp", "pulsa", "aku", "saya"}:
            return m.group(1)
        m = re.search(r"\b(?:beliin|isiin|isi|beli)\s+([a-z]+)\s+pulsa", lowered)
        if m and m.group(1) not in {"pulsa", "nomor", "nomer"}:
            return m.group(1)
        return None

    @staticmethod
    def _parse_amount(lowered: str) -> Optional[int]:
        for slang, value in _SLANG_AMOUNTS.items():
            if slang in lowered:
                return value
        m = re.search(r"(\d+(?:[.,]\d+)?)\s*(rb|ribu|k|jt|juta)\b", lowered)
        if m:
            base = float(m.group(1).replace(",", "."))
            mult = 1_000 if m.group(2) in {"rb", "ribu", "k"} else 1_000_000
            return int(base * mult)
        m = re.search(r"\b(\d{4,9})\b(?!\s*(?:rb|ribu|k|jt|juta))", lowered)
        if m and not m.group(1).startswith("08"):
            value = int(m.group(1))
            if 500 <= value <= 100_000_000:
                return value
        return None


def create_engine(offline: bool = False) -> "BaseNLUEngine":
    """Factory: live Gemini engine, or mock when offline / no key present."""
    if offline:
        return MockNLUEngine()
    try:
        return NLUEngine()
    except ValueError:
        logger.warning("No GEMINI_API_KEY found — falling back to MockNLUEngine.")
        return MockNLUEngine()