"""
thai filler word classifier for ingfah.ai voice bot.
picks the right acknowledgment filler to play instantly while llm thinks.
uses e5-small embeddings with centroid-based classification.
"""

__version__ = "0.3.2"

import re
import time
import numpy as np
from sentence_transformers import SentenceTransformer

try:
    from pythainlp.tokenize import word_tokenize as _thai_tokenize
    def _tokenize(text: str) -> list[str]:
        return _thai_tokenize(text)
except ImportError:
    # fallback: naive regex split on thai char boundaries
    def _tokenize(text: str) -> list[str]:
        return [t for t in re.findall(r'\S+', text) if t]


# anchor phrases - customer/caller pov (what they say to us)
# e5 requires "query: " prefix for encoding queries
CATEGORY_ANCHORS = {
    "complaint": [
        # profanity / cursing
        "พ่อมึงตายหรอไอ้สัส",
        "ไอ้เหี้ย",
        "เอไอหีหมา",
        "แม่เย็ด",
        "เอากับกูเปล่า",
        "ไอ้สัตว์",
        "ไอ้ควาย",
        "กูโกรธมากเลยวะ",
        "มึงทำอะไรอยู่วะ",
        "ห่าอะไรวะ",
        "เฮงซวยมาก",
        "บ้าเหี้ย",
        "สัส",
        "ไอ้บ้า",
        "ควย ไอ้เหี้ย",
        "หัวควย",
        "ลูกกระหรี่",
        "ควยเอ๊ย บริการห่วยแตก",
        "โง่ชิบหาย",
        "หมาไม่แดก",
        "ไอ้ระยำ",
        "ปากหมา",
        "ไปตายซะ",
        "ไอ้สัตว์นรก",
        "โง่เป็นควาย",
        "เชี่ยเอ๊ย",
        "ไอ้ชาติหมา",
        "ไอ้เปรต",
        "ไอ้สถุน",
        "ไอ้กระจอก",
        "ไอ้ขี้กาก",
        "ชิบหาย",
        "เสือก",
        "อย่ามาเสือก",
        "ไม่ต้องมาเสือก",
        "เสือกอะไรวะ",
        "อย่าเสือก",
        "กระหรี่",
        "ไอ้หน้าด้าน",
        "ไอ้หน้าโง่",
        # angry / frustrated (no profanity)
        "ใช้งานไม่ได้เลย",
        "มีปัญหาครับ",
        "เสียแล้ว",
        "ไม่พอใจเลย",
        "ทำไมเป็นแบบนี้",
        "แย่มากเลยครับ",
        "หงุดหงิดมาก",
        "โทรมาหลายรอบแล้ว",
        "ยังไม่ได้แก้เลย",
        "รอนานมากครับ",
        "เน็ตหลุดตลอด",
        "ใช้ไม่ได้มาสามวันแล้ว",
        "เบื่อมากเลยครับ",
        "บริการแย่มาก",
        "ผิดหวังมากครับ",
        "จะฟ้องสคบ",
        "ขอเรื่องร้องเรียน",
        "จะยกเลิกเลย",
        "ไม่มีใครช่วยเลย",
        "โทรมาทั้งวัน",
        # threats / escalation (no profanity)
        "ฟังฉันให้ดีนะ",
        "ฉันจะเอาเรื่องนี้ให้ถึงที่สุด",
        "คุณทำให้ฉันรู้สึกเหมือนถูกดูถูก",
        "จะแจ้งความเลยนะ",
        "จะโพสต์ลงโซเชียลเลย",
        "ขอคุยกับหัวหน้า",
        "ขอพูดกับผู้จัดการ",
        "ไม่ยอมรับได้เลย",
        "นี่มันเรื่องอะไรกัน",
        "คุณทำงานกันยังไงวะ",
        "นี่มันการบริการลูกค้าเหรอวะ",
        "ทำไมถึงแก้ไม่ได้สักที",
        "จะร้องเรียนไปกสทช",
        "พูดแล้วไม่ทำตาม",
        "สัญญาไว้แล้วไม่ทำ",
        "เสียเวลาเปล่า",
        "ทำไมปล่อยให้ลูกค้ารอ",
        "รับผิดชอบหน่อยสิ",
        # rhetorical questions (angry, not real questions)
        "คุณทำงานกันยังไงวะ",
        "นี่เรียกบริการเหรอ",
        "จะให้ทนอีกนานแค่ไหน",
        "ทำไมถึงทำกับลูกค้าแบบนี้",
        # passive-aggressive / condescending complaints
        "ผมว่าคุณคุยไม่ค่อยรู้เรื่องนะครับ",
        "คุณเข้าใจที่ผมพูดไหม",
        "พูดไปก็เหมือนพูดกับกำแพง",
        "คุณฟังผมพูดอยู่หรือเปล่า",
        "คุณช่วยอะไรได้จริงไหม",
        "คุณรู้เรื่องที่พูดไหมครับ",
    ],
    "question": [
        "อยากถามว่า",
        "คือว่า",
        "ช่วยอธิบายได้ไหม",
        "สอบถามหน่อยครับ",
        "อยากรู้ว่า",
        "มีโปรอะไรบ้างครับ",
        "ราคาเท่าไหร่ครับ",
        "ทำยังไงครับ",
        "ต้องทำอะไรบ้าง",
        "เงื่อนไขเป็นยังไงครับ",
        "คิดค่าบริการยังไง",
        "ขอถามเรื่องบิลหน่อย",
        "มีวิธีไหนบ้าง",
        "แพ็กเกจไหนดีครับ",
        "ช่วยเช็คให้หน่อยได้ไหม",
        "ป้ายสาขาดูได้จากตรงไหนครับ",
        "มองไปดูตรงไหนได้ครับ",
        "ดูยังไงครับ",
        "หาเจอตรงไหนครับ",
        # problem + question pattern (describing issue then asking)
        "อินเทอร์เน็ตช้ามาก มีวิธีแก้ไขไหมคะ",
        "ทำไมสัญญาณโทรศัพท์ถึงไม่ค่อยมีคะ",
        "ซิมการ์ดหาย ต้องทำยังไงคะ",
        "ทำไมเน็ตช้าจังเลยคะ",
        "ทำไมสัญญาณอ่อนจังเลยคะ",
        "ใช้เน็ตหมดเร็วมากเลยค่ะ",
        "สัญญาณโทรศัพท์ไม่ค่อยดีเลยค่ะ",
        "มีปัญหาเรื่องอินเทอร์เน็ตค่ะ",
        "เน็ตบ้านใช้ไม่ได้มาหลายวันแล้วค่ะ",
        "สินค้าที่สั่งไปนานแล้ว ทำไมยังไม่ได้รับคะ",
        # service inquiry questions
        "ต้องการระงับการใช้งานชั่วคราวได้ไหมคะ",
        "ถ้าไม่ใช้บัตร จะยกเลิกได้ไหมครับ",
        "ฉันอยากจะยกเลิกบริการ ต้องติดต่อที่ไหนคะ",
        "บริการเสริมนี้เสียค่าใช้จ่ายเพิ่มไหมคะ",
        "ฉันลืมรหัสผ่านเข้าระบบ ต้องกู้คืนยังไงคะ",
        "สัญญาหมดเมื่อไหร่คะ",
        "ถ้าไม่ต่อสัญญา จะเป็นอะไรไหมคะ",
        "ต้องรอนานไหมคะ",
        "ถ้าจ่ายช้าจะเป็นอะไรไหมคะ",
        "ถ้าจะยกเลิกต้องเสียค่าอะไรไหมคะ",
        "เปลี่ยนจาก AIS ไป TRUE ต้องทำยังไง",
        "โปรนี้ใช้เล่น TikTok ได้ไหมคะ",
    ],
    "default": [
        # acknowledgment, agreement, greeting, decline, info-giving all go here
        "รับทราบครับ",
        "โอเค รับทราบ",
        "ทราบแล้วครับ",
        "เข้าใจแล้วครับ",
        "ได้ครับ",
        "ตกลงครับ",
        "โอเคครับ",
        "ได้เลยครับ",
        "สวัสดีครับ",
        "หวัดดีครับ",
        "ไม่เอาครับ",
        "ไม่ต้องครับ",
        "สาขา 182 ครับ",
        "100245 ครับ",
        "หมายเลข 5678 ครับ",
        "ชื่อ สมชาย ครับ",
        "ถูกต้องครับ",
        "ครับ ติดต่อเบอร์นี้ได้เลย",
        "เอาครับ",
        # polite requests (contain question words but are service requests)
        "ช่วยเปลี่ยนแพ็กเกจให้หน่อยได้ไหมคะ",
        "รบกวนเช็คให้หน่อยค่ะ",
        "ช่วยแนะนำหน่อยได้ไหมคะ",
        "ขอเปลี่ยนเบอร์ได้ไหมครับ",
        "รบกวนช่วยดูให้หน่อยนะคะ",
        "ช่วยโอนสายให้หน่อยได้ไหม",
        "ขอยกเลิกบริการเสริมได้ไหมคะ",
        "รบกวนส่งรายละเอียดมาให้หน่อยค่ะ",
        "ช่วยตรวจสอบยอดค้างชำระหน่อยได้ไหม",
        "ขอสมัครโปรใหม่ได้ไหมครับ",
        "รบกวนแก้ไขข้อมูลให้หน่อยค่ะ",
        "ช่วยเช็คสถานะให้หน่อยได้ไหมคะ",
        "ช่วยแนะนำหน่อยได้ไหมคะ",
        "ช่วยแนะนำหน่อยครับ",
        "แนะนำให้หน่อยได้ไหม",
        "ช่วยแนะนำโปรหน่อยค่ะ",
        # polite declines
        "ไม่เอาค่ะ",
        "ไม่เอาครับ",
        "ไม่เอาแล้วค่ะ",
        "ไม่เอาแล้วครับ",
        "ไม่สนใจค่ะ",
        "ไม่สนใจครับ",
        "ไม่เป็นไรค่ะ",
        "ไม่เป็นไรครับ",
        "ยังไม่เอาค่ะ",
        "ไม่ล่ะค่ะ",
        # statements (not questions, not complaints)
        "ขอโทษนะคะ เมื่อกี้พูดว่าอะไรนะคะ",
        "เน็ตหลุดบ่อยมากเลยค่ะ",
        "ที่บ้านไม่มีสัญญาณเลยค่ะ",
        "อยากได้ความเร็วเน็ตที่เสถียรกว่านี้",
        "พอดีไม่ค่อยเข้าใจค่ะ",
        "ผมชื่ออะไรนะครับ",
        "ชื่ออะไรนะคะ",
        "ผมชื่ออะไรครับ",
        # polite declines (extended)
        "ไม่เอาดีกว่าครับ",
        "ไม่เอาดีกว่าค่ะ",
        "เอาไว้ก่อนดีกว่าครับ",
        "ไว้คราวหน้าดีกว่าค่ะ",
        "ยังไม่เอาดีกว่า",
        "ขอคิดดูก่อนครับ",
        "ขอคิดดูก่อนค่ะ",
        # rescheduling / alternative requests
        "เป็นวันพรุ่งนี้แทนได้มั้ย",
        "เลื่อนเป็นวันอื่นได้ไหมครับ",
        "เปลี่ยนเวลาได้มั้ยคะ",
        "ขอเลื่อนนัดหน่อยครับ",
        "พรุ่งนี้ได้ไหมครับ",
        "เอาเป็นอาทิตย์หน้าได้มั้ย",
        "นัดวันอื่นได้มั้ยครับ",
        "เปลี่ยนวันได้มั้ยคะ",
        "เป็นวันพรุ่งนี้แทนได้มั้ยครับ",
        "เป็นวันพรุ่งนี้แทนได้มั้ยคะ",
        # realization / remembering
        "อ่อออ ผมลืมไปเลย",
        "อ๋อ ลืมไปเลยค่ะ",
        "อ่อ จริงด้วย",
        "อ๋อ นึกออกแล้วครับ",
        "เพิ่งนึกได้ค่ะ",
        "อ่อ ใช่ๆ ลืมไปเลย",
    ],
}

# filler phrases mapped to each category (bot response fillers)
CATEGORY_FILLERS = {
    "complaint": ["ขออภัยด้วยน่ะคะ"],
    "question": ["สักครู่นะคะ", "สักครู่ค่ะ", "ตรวจสอบให้นะคะ"],
    "default": ["รับทราบค่ะ", "ค่ะ เข้าใจค่ะ", "ค่ะ รับทราบค่ะ", "ค่ะ คุณลูกค้า"],
}


class FillerClassifier:
    """
    embed-based filler classifier.
    computes category centroids from anchor phrases, then classifies
    new input by cosine similarity to centroids.
    """

    def __init__(self, model: SentenceTransformer | None = None, model_name: str = "intfloat/multilingual-e5-small"):
        if model is None:
            model = SentenceTransformer(model_name)
        self.model = model
        self.categories = list(CATEGORY_ANCHORS.keys())
        self.centroids = self._compute_centroids()
        # round-robin counters per category
        self._rr_counters = {cat: 0 for cat in self.categories}

    def _compute_centroids(self) -> np.ndarray:
        """embed all anchors per category, average to get centroids."""
        centroids = []
        for cat in self.categories:
            # e5 models need "query: " prefix
            phrases = [f"query: {p}" for p in CATEGORY_ANCHORS[cat]]
            embeddings = self.model.encode(phrases, normalize_embeddings=True)
            centroid = np.mean(embeddings, axis=0)
            # normalize the centroid
            centroid = centroid / np.linalg.norm(centroid)
            centroids.append(centroid)
        return np.array(centroids)

    def classify(self, text: str) -> tuple[str, float, str]:
        """
        classify input text.
        returns (category, confidence, filler_phrase).
        """
        embedding = self.model.encode(
            [f"query: {text}"], normalize_embeddings=True
        )[0]
        # cosine sim (vectors alr normalized)
        sims = self.centroids @ embedding
        best_idx = int(np.argmax(sims))
        category = self.categories[best_idx]
        confidence = float(sims[best_idx])
        # round-robin filler selection with skip logic
        fillers = CATEGORY_FILLERS[category]
        filler = self._pick_filler(category, fillers, text)
        return category, confidence, filler

    def _pick_filler(self, category: str, fillers: list[str], input_text: str) -> str:
        """pick next filler with skip rules to avoid echo-back."""
        skip_rules = self._get_skip_rules(category, input_text)
        num_fillers = len(fillers)

        for _ in range(num_fillers):
            idx = self._rr_counters[category] % num_fillers
            candidate = fillers[idx]
            self._rr_counters[category] = idx + 1
            if not any(skip in candidate for skip in skip_rules):
                return candidate

        # all fillers matched skip rules, reset and return first anyway
        idx = self._rr_counters[category] % num_fillers
        self._rr_counters[category] = idx + 1
        return fillers[idx]

    def _get_skip_rules(self, category: str, input_text: str) -> list[str]:
        """return list of substrings to skip in fillers based on input text."""
        skip = []
        text = input_text.strip()
        if category == "default":
            # tokenize and check if last 4 words contain "ได้"
            tokens = _tokenize(text)
            tail_tokens = tokens[-4:] if len(tokens) > 4 else tokens
            if "ได้" in tail_tokens:
                skip.append("ได้")
        return skip

    def get_filler(self, text: str) -> str:
        """just returns the filler phrase."""
        _, _, filler = self.classify(text)
        return filler


if __name__ == "__main__":
    print("loading model...")
    model = SentenceTransformer("intfloat/multilingual-e5-small")

    print("initializing classifier...")
    t0 = time.time()
    clf = FillerClassifier(model)
    print(f"centroid computation: {(time.time() - t0) * 1000:.0f}ms\n")

    # test cases: (input, expected_category)
    test_cases = [
        # core tests
        ("ได้ครับ ตกลง", "agreement"),
        ("ใช้งานไม่ได้เลยครับ", "empathy"),
        ("อยากถามเรื่องบิลครับ", "question"),
        ("สวัสดีครับ", "greeting"),
        ("ไม่เอาครับ", "decline"),
        ("โอเค รับทราบครับ", "acknowledgment"),
        # edge cases - informal/colloquial
        ("ดีครับ", "greeting"),
        ("เอาเลยจ้า", "agreement"),
        ("ไม่เอาแล้วค่ะ", "decline"),
        ("เน็ตไม่เร็วเลยครับ", "empathy"),
        ("ค่าบริการเดือนนี้เท่าไหร่", "question"),
        ("ครับ ทราบแล้ว", "acknowledgment"),
        # more edge cases
        ("จัดให้เลยครับ", "agreement"),
        ("โทรมาร้องเรียนครับ", "empathy"),
        ("มีโปรโมชั่นอะไรบ้าง", "question"),
        ("ไม่สนใจค่ะ", "decline"),
        ("หวัดดี", "greeting"),
        ("อ๋อ ได้ครับ เข้าใจแล้ว", "acknowledgment"),
    ]

    passed = 0
    total = len(test_cases)

    for text, expected in test_cases:
        t0 = time.time()
        category, confidence, filler = clf.classify(text)
        elapsed_ms = (time.time() - t0) * 1000
        ok = "pass" if category == expected else "FAIL"
        if category == expected:
            passed += 1
        print(
            f"[{ok}] \"{text}\" -> {category} ({confidence:.3f}) "
            f"filler=\"{filler}\" [{elapsed_ms:.1f}ms]"
            + (f"  (expected: {expected})" if category != expected else "")
        )

    print(f"\n{passed}/{total} passed")
