"""
Postprocessing utilities for BISINDO Trans.

Functions for cleaning and formatting model output.
"""

import re
from typing import Dict


# Default phrase corrections
DEFAULT_PHRASE_CORRECTIONS: Dict[str, str] = {
    "cara yang belajar": "cara belajar",
    "mengapa itu": "mengapa",
    "mana yang di": "di mana",
    "mana yang kamu": "dari mana kamu",
    "siapa itu nama": "siapa nama",
    "apa itu kabar": "apa kabar",
    "akan lagi akan": "sedang",
    "akan suka": "suka",
    "nama yang": "nama",
    "muhamadaldi": "Muhammad Aldi",
    "  ": " ",
}


def bersihkan_output(raw_text: str, phrase_corrections: Dict[str, str] = None) -> str:
    """
    Membersihkan dan memformat output dari model.
    
    Steps:
        1. Bersihkan tag __NAME__
        2. Ganti frasa yang salah
        3. Format kapitalisasi dan punctuation
    
    Args:
        raw_text: Raw output dari model
        phrase_corrections: Optional custom phrase corrections dict
        
    Returns:
        String yang sudah dibersihkan dan diformat
        
    Example:
        >>> bersihkan_output("saya __NAME__aldi makan")
        'Saya Aldi makan.'
    """
    if phrase_corrections is None:
        phrase_corrections = DEFAULT_PHRASE_CORRECTIONS
    
    # 1. Bersihkan Nama (__NAME__)
    words = raw_text.split()
    clean_words = []
    for w in words:
        if re.search(r'__NAME__', w, re.IGNORECASE):
            name_clean = re.sub(r'__NAME__', '', w, flags=re.IGNORECASE)
            clean_words.append(name_clean.title())
        else:
            clean_words.append(w.replace("_", " "))
    
    sentence = " ".join(clean_words)
    
    # 2. PHRASE CLEANING (Kamus Koreksi)
    sentence_lower = sentence.lower()
    for bad, good in phrase_corrections.items():
        if bad in sentence_lower:
            pattern = re.compile(re.escape(bad), re.IGNORECASE)
            sentence = pattern.sub(good, sentence)
    
    # 3. Deteksi Kalimat Tanya
    kata_tanya = ["apa", "siapa", "kapan", "dimana", "kemana", "darimana", 
                  "bagaimana", "mengapa", "kenapa", "berapa", "mana"]
    is_question = False
    sentence_lower = sentence.lower()
    
    # Cek apakah kalimat dimulai dengan kata tanya atau mengandung kata tanya
    words_lower = sentence_lower.split()
    if words_lower:
        # Kalimat dimulai dengan kata tanya
        if words_lower[0] in kata_tanya:
            is_question = True
        # Kalimat mengandung "ke mana", "di mana", "dari mana"
        if "ke mana" in sentence_lower or "di mana" in sentence_lower or "dari mana" in sentence_lower:
            is_question = True
        # Kalimat mengandung kata tanya di posisi manapun
        for kt in kata_tanya:
            if kt in words_lower:
                is_question = True
                break
    
    # 4. Penempatan Koma
    # Aturan: setelah kata penghubung di awal kalimat, sebelum konjungsi
    kata_perlu_koma_setelah = ["namun", "tetapi", "akan tetapi", "oleh karena itu", 
                               "selain itu", "dengan demikian", "jadi"]
    kata_perlu_koma_sebelum = ["tetapi", "namun", "sedangkan", "melainkan"]
    
    words = sentence.split()
    new_words = []
    for i, word in enumerate(words):
        word_lower = word.lower().rstrip(".,!?")
        
        # Koma sebelum konjungsi (jika bukan di awal)
        if word_lower in kata_perlu_koma_sebelum and i > 0:
            # Tambahkan koma ke kata sebelumnya jika belum ada
            if new_words and not new_words[-1].endswith(","):
                new_words[-1] = new_words[-1] + ","
        
        new_words.append(word)
        
        # Koma setelah kata penghubung di awal
        if i == 0 and word_lower in kata_perlu_koma_setelah:
            if not word.endswith(","):
                new_words[-1] = new_words[-1] + ","
    
    sentence = " ".join(new_words)
    
    # 5. Final Formatting
    sentence = re.sub(r'\s+', ' ', sentence).strip()
    if sentence:
        final_output = sentence[0].upper() + sentence[1:]
    else:
        final_output = ""
    
    # 6. Tanda Baca Akhir
    if not final_output.endswith(("?", "!", ".")):
        if is_question:
            final_output += "?"
        else:
            final_output += "."

    final_output = final_output.replace(" ?", "?").replace(" .", ".").replace(" ,", ",")
    return final_output


# --- TEST ---
if __name__ == "__main__":
    tests = [
        "saya __NAME__aldi makan",
        "cara yang belajar matematika",
        "apa itu kabar kamu",
    ]
    
    for t in tests:
        print(f"Input : {t}")
        print(f"Output: {bersihkan_output(t)}")
        print()
