"""
NLU Module — Multi-Intent Decomposition + Entity Extraction
=============================================================
Fixes P0/P1 from the feedback:

  - Decomposes ONE user message into a LIST of tasks (compound requests)
  - Extracts ALL entities present in the message (slot prefill —
    "send money to abu" never re-asks for the recipient)
  - Returns per-task confidence so destructive intents can be gated
  - Distinguishes "ask about X" from "do X" (branch location ≠ block card)

Backend chain (first available wins):
  1. LLM_API    — HF Serverless Inference (set HF_TOKEN) — best quality
  2. LLM_LOCAL  — Qwen2.5-1.5B-Instruct loaded in-process — good, slower
  3. RULES      — improved keyword rules — degraded but never crashes

All backends return the same schema:

  {
    "tasks": [
      {
        "intent": "send_money",
        "confidence": 0.93,
        "slots": {"recipient": "abu", "amount": "350000"},
        "utterance_span": "send 350000 to abu"
      },
      ...
    ],
    "backend": "llm_api"
  }
"""

import os
import re
import json
import logging
from typing import Optional

logger = logging.getLogger(__name__)

# ── Intent catalogue (shared by all backends) ────────────────────────────────
INTENT_SCHEMA = {
    "greeting":         {"slots": [],                          "destructive": False},
    "balance_inquiry":  {"slots": ["account_id"],              "destructive": False},
    "send_money":       {"slots": ["recipient", "amount", "account_id"], "destructive": True},
    "bill_payment":     {"slots": ["account_id", "amount"],    "destructive": True},
    "block_card":       {"slots": ["account_id"],              "destructive": True},
    "branch_info":      {"slots": [],                          "destructive": False},
    "card_request":     {"slots": [],                          "destructive": False},
    "report_issue":     {"slots": ["issue_desc"],              "destructive": False},
    "track_order":      {"slots": ["order_id"],                "destructive": False},
    "return_item":      {"slots": ["order_id", "return_reason"], "destructive": False},
    "human_agent":      {"slots": [],                          "destructive": False},
    "confirmation_yes": {"slots": [],                          "destructive": False},
    "confirmation_no":  {"slots": [],                          "destructive": False},
    "cancel":           {"slots": [],                          "destructive": False},
    "goodbye":          {"slots": [],                          "destructive": False},
    "unknown":          {"slots": [],                          "destructive": False},
}

NLU_SYSTEM_PROMPT = """You are the NLU module of a customer-service voice agent.
Decompose the user's message into ALL tasks it contains, in order.
Extract every entity present. Never invent entities that are not in the text.

Intents: greeting, balance_inquiry, send_money, bill_payment, block_card,
branch_info, card_request, report_issue, track_order, return_item,
human_agent, confirmation_yes, confirmation_no, cancel, goodbye, unknown.

Slots: recipient, amount, account_id, location, issue_desc, order_id, return_reason.

CRITICAL disambiguation rules:
- "where is your branch so I can get my card" = branch_info + card_request.
  It is NOT block_card. Only choose block_card if the user explicitly wants to
  BLOCK, FREEZE, or DEACTIVATE a card.
- A message can contain multiple tasks joined by "and", "also", "then".
  Output one task per action. "check my balance and send 5000 to musa"
  = [balance_inquiry, send_money{recipient: musa, amount: 5000}].
- If the user answers a question (e.g. gives a reason like "too small"),
  map it to the slot of the pending task, intent = the pending intent.
- Confidence in [0,1]: how sure you are of the INTENT (not the slots).

Respond with ONLY valid JSON, no markdown, no commentary:
{"tasks":[{"intent":"...","confidence":0.0,"slots":{},"utterance_span":"..."}]}"""


class NLU:

    def __init__(self, prefer: str = "auto"):
        self.hf_token = os.getenv("HF_TOKEN", "")
        self.api_model = os.getenv(
            "NLU_API_MODEL", "Qwen/Qwen2.5-72B-Instruct")
        self.local_model_id = os.getenv(
            "NLU_LOCAL_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
        # The local backend is OPT-IN. Left automatic, the first NLU call on a
        # Space silently downloads a ~3GB model mid-demo and blocks for
        # minutes. Enable deliberately with NLU_LOCAL=1 on hardware that can
        # take it.
        self.use_local = os.getenv("NLU_LOCAL", "0") == "1"
        self._local_pipe = None
        self._dead = set()          # backends that failed — never retried
        self.prefer = prefer

    # ── Public API ────────────────────────────────────────────────────────────

    def parse(self, text: str, pending_intent: Optional[str] = None,
              pending_slot: Optional[str] = None) -> dict:
        """
        text           : English pivot text of the user turn
        pending_intent : intent currently awaiting a slot (context for the LLM)
        pending_slot   : which slot we asked for last turn
        """
        context = ""
        if pending_intent and pending_slot:
            context = (f"\nContext: you previously asked the user for the "
                       f"'{pending_slot}' of a '{pending_intent}' task. "
                       f"A short answer likely fills that slot.")

        for backend in self._backend_order():
            name = backend.__name__
            if name in self._dead:
                continue
            try:
                result = backend(text, context)
                if result and result.get("tasks"):
                    result = self._sanitize(result)
                    logger.info(f"NLU[{result['backend']}]: "
                                f"{json.dumps(result['tasks'])[:200]}")
                    return result
            except Exception as e:
                # Mark dead so a missing dependency or bad token doesn't cost
                # a retry (and a re-download attempt) on every single turn.
                self._dead.add(name)
                logger.warning(
                    f"NLU backend {name} failed and is disabled for this "
                    f"session: {e}")
        # Absolute last resort
        return {"tasks": [{"intent": "unknown", "confidence": 0.0,
                            "slots": {}, "utterance_span": text}],
                "backend": "none"}

    # ── Backend chain ─────────────────────────────────────────────────────────

    def _backend_order(self):
        if self.prefer == "rules":
            return [self._rules_backend]
        chain = []
        if self.hf_token:
            chain.append(self._api_backend)
        if self.use_local:
            chain.append(self._local_backend)
        chain.append(self._rules_backend)
        return chain

    # ── 1. HF Serverless Inference API ───────────────────────────────────────

    def _api_backend(self, text: str, context: str) -> Optional[dict]:
        import requests
        url = f"https://api-inference.huggingface.co/models/{self.api_model}/v1/chat/completions"
        payload = {
            "model": self.api_model,
            "messages": [
                {"role": "system", "content": NLU_SYSTEM_PROMPT + context},
                {"role": "user",   "content": text},
            ],
            "max_tokens": 400,
            "temperature": 0.1,
        }
        r = requests.post(url, json=payload, timeout=20,
                          headers={"Authorization": f"Bearer {self.hf_token}"})
        r.raise_for_status()
        raw = r.json()["choices"][0]["message"]["content"]
        parsed = self._extract_json(raw)
        if parsed:
            parsed["backend"] = "llm_api"
        return parsed

    # ── 2. Local small LLM ────────────────────────────────────────────────────

    def _local_backend(self, text: str, context: str) -> Optional[dict]:
        if self._local_pipe is None:
            logger.info(f"Loading local NLU model {self.local_model_id} …")
            from transformers import pipeline as hf_pipeline
            import torch
            self._local_pipe = hf_pipeline(
                "text-generation",
                model=self.local_model_id,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                device_map="auto",
            )
        messages = [
            {"role": "system", "content": NLU_SYSTEM_PROMPT + context},
            {"role": "user",   "content": text},
        ]
        out = self._local_pipe(messages, max_new_tokens=400,
                                do_sample=False, temperature=None, top_p=None)
        raw = out[0]["generated_text"][-1]["content"]
        parsed = self._extract_json(raw)
        if parsed:
            parsed["backend"] = "llm_local"
        return parsed

    # ── 3. Improved rules (never fails) ───────────────────────────────────────

    def _rules_backend(self, text: str, context: str) -> dict:
        """
        Better than the old FSM keywords:
          - splits on conjunctions to find MULTIPLE tasks
          - extracts entities per clause
          - branch_info vs block_card disambiguation
        """
        t = text.lower().strip()

        # Split compound message into clauses
        clauses = re.split(r'\b(?:and also|and then|then|and|also|;|\. )\b', t)
        clauses = [c.strip() for c in clauses if c.strip()]

        tasks = []
        for clause in clauses:
            task = self._rules_classify_clause(clause)
            if task:
                tasks.append(task)

        # Merge duplicate consecutive intents (e.g. "and" split an entity off)
        merged = []
        for task in tasks:
            if merged and merged[-1]["intent"] == task["intent"]:
                merged[-1]["slots"].update(task["slots"])
                merged[-1]["utterance_span"] += " " + task["utterance_span"]
            else:
                merged.append(task)

        if not merged:
            merged = [{"intent": "unknown", "confidence": 0.3,
                       "slots": {}, "utterance_span": t}]

        return {"tasks": merged, "backend": "rules"}

    def _rules_classify_clause(self, clause: str) -> Optional[dict]:
        slots = {}

        # ── Entity extraction (always, regardless of intent) ────────────────
        # In a money-action clause ("send/transfer/pay X to Y"), the number is
        # an AMOUNT. Only treat 6-12 digit numbers as account_id when the
        # clause is about the account itself, or there is no money verb.
        money_verb = any(v in clause for v in ("send", "transfer", "pay"))
        account_ctx = any(v in clause for v in ("account", "acct", "number is"))
        numbers = re.findall(r'\b\d[\d,\.]*\b', clause)
        for num in numbers:
            digits = num.replace(",", "").replace(".", "")
            if money_verb and "amount" not in slots and len(digits) <= 7:
                slots["amount"] = digits
            elif (account_ctx or not money_verb) and 6 <= len(digits) <= 12 \
                    and "account_id" not in slots:
                slots["account_id"] = digits
            elif "amount" not in slots and len(digits) <= 7:
                slots["amount"] = digits
        # recipient: "to <name>" — take the LAST valid match, skipping verbs
        # ("I want to send money to abu" must yield 'abu', not 'send')
        RECIPIENT_STOPWORDS = {
            "my", "the", "a", "an", "me", "you", "check", "send", "transfer",
            "pay", "get", "make", "do", "know", "see", "block", "return",
            "track", "him", "her", "them", "it", "confirm", "cancel"}
        for m in re.finditer(r'\bto\s+([a-z]{2,20})\b', clause):
            name = m.group(1)
            if name not in RECIPIENT_STOPWORDS:
                slots["recipient"] = name
        # order id
        m = re.search(r'\border\s*#?\s*([a-z0-9\-]{4,20})\b', clause)
        if m:
            slots["order_id"] = m.group(1)

        # ── Intent (order matters: destructive intents need explicit verbs) ──
        def has(*kws):
            return any(kw in clause for kw in kws)

        # branch/location questions BEFORE block_card — fixes P0 #2
        if has("branch", "closest", "nearest", "location", "where is", "address"):
            intent, conf = "branch_info", 0.85
            if has("card", "atm"):
                # compound: they also want a card — but NOT to block it
                return {"intent": "branch_info", "confidence": 0.85,
                        "slots": slots, "utterance_span": clause}
        elif has("block my card", "block card", "freeze", "deactivate", "stolen", "lost my card"):
            intent, conf = "block_card", 0.8
        elif has("send", "transfer") and (slots.get("recipient") or slots.get("amount")):
            intent, conf = "send_money", 0.85
        elif has("send money", "transfer money"):
            intent, conf = "send_money", 0.75
        elif has("balance", "how much", "asusun"):
            intent, conf = "balance_inquiry", 0.85
        elif has("pay", "bill", "recharge", "invoice"):
            intent, conf = "bill_payment", 0.75
        elif has("track", "where is my order", "delivery", "shipment"):
            intent, conf = "track_order", 0.8
        elif has("return", "refund", "send back"):
            intent, conf = "return_item", 0.8
        elif has("problem", "issue", "complaint", "not working", "error"):
            intent, conf = "report_issue", 0.7
            slots["issue_desc"] = clause
        elif has("human", "agent", "person", "operator", "representative"):
            intent, conf = "human_agent", 0.9
        elif has("yes", "yep", "correct", "confirm", "sure", "okay", "ok"):
            intent, conf = "confirmation_yes", 0.8
        elif has("no", "nope", "wrong", "cancel that"):
            intent, conf = "confirmation_no", 0.8
        elif has("hello", "hi ", "good morning", "sannu", "salam"):
            intent, conf = "greeting", 0.9
        elif has("bye", "goodbye", "thank"):
            intent, conf = "goodbye", 0.85
        else:
            return {"intent": "unknown", "confidence": 0.3,
                    "slots": slots, "utterance_span": clause}

        return {"intent": intent, "confidence": conf,
                "slots": slots, "utterance_span": clause}

    # ── Helpers ───────────────────────────────────────────────────────────────

    @staticmethod
    def _extract_json(raw: str) -> Optional[dict]:
        """Robustly pull the first JSON object out of LLM output."""
        raw = raw.strip()
        raw = re.sub(r'^```(?:json)?|```$', '', raw, flags=re.MULTILINE).strip()
        # find first { … matching last }
        start = raw.find("{")
        end   = raw.rfind("}")
        if start == -1 or end == -1:
            return None
        try:
            return json.loads(raw[start:end + 1])
        except json.JSONDecodeError:
            return None

    @staticmethod
    def _sanitize(result: dict) -> dict:
        """Validate schema, clamp confidence, drop hallucinated slots."""
        valid_slots = {"recipient", "amount", "account_id", "location",
                       "issue_desc", "order_id", "return_reason"}
        clean_tasks = []
        for task in result.get("tasks", []):
            intent = task.get("intent", "unknown")
            if intent not in INTENT_SCHEMA:
                intent = "unknown"
            conf = float(task.get("confidence", 0.5))
            conf = max(0.0, min(1.0, conf))
            slots = {k: str(v).strip() for k, v in (task.get("slots") or {}).items()
                     if k in valid_slots and v not in (None, "", "null", "None")}
            clean_tasks.append({
                "intent": intent, "confidence": conf, "slots": slots,
                "utterance_span": str(task.get("utterance_span", ""))[:200],
            })
        result["tasks"] = clean_tasks or [
            {"intent": "unknown", "confidence": 0.0, "slots": {},
             "utterance_span": ""}]
        return result