"""
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. Final Formatting
    sentence = re.sub(r'\s+', ' ', sentence).strip()
    if sentence:
        final_output = sentence[0].upper() + sentence[1:]
    else:
        final_output = ""
        
    if not final_output.endswith(("?", "!", ".")):
        final_output += "."

    final_output = final_output.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()
