src/services/query_normalizer.py

from __future__ import annotations

import logging
import re
import unicodedata
from dataclasses import dataclass

logger = logging.getLogger(__name__)

ABBREVIATIONS: dict[str, str] = {
    "ND": "Nghi dinh",
    "NĐ": "Nghị định",
    "TT": "Thong tu",
    "BL": "Bo luat",
    "CP": "Chinh phu",
    "QH": "Quoc hoi",
    "UBND": "Uy ban nhan dan",
    "HDND": "Hoi dong nhan dan",
    "HĐND": "Hội đồng nhân dân",
    "BLDS": "Bo luat Dan su",
    "BLHS": "Bo luat Hinh su",
}

ARTICLE_PATTERN = re.compile(r"\b(?:dieu|điều)\s+(\d+[a-z]?)\b", re.IGNORECASE)
CLAUSE_PATTERN = re.compile(r"\b(?:khoan|khoản)\s+(\d+[a-z]?)\b", re.IGNORECASE)
POINT_PATTERN = re.compile(r"\b(?:diem|điểm)\s+([a-z])\b", re.IGNORECASE)


@dataclass(frozen=True)
class NormalizedQuery:
    normalized_text: str
    extracted_refs: dict[str, str]
    abbreviations_expanded: list[str]
    original_query: str


def normalize_query(query: str) -> NormalizedQuery:
    """Normalize Vietnamese legal query text for matching and reference extraction."""
    original = str(query or "")
    expanded_text, abbreviations = _expand_abbreviations(original)
    normalized_text = _normalize_for_matching(expanded_text)
    refs = _extract_refs(normalized_text)
    logger.info(
        "Normalized query; abbreviations=%s refs=%s",
        abbreviations,
        refs,
    )
    return NormalizedQuery(
        normalized_text=normalized_text,
        extracted_refs=refs,
        abbreviations_expanded=abbreviations,
        original_query=original,
    )


def _expand_abbreviations(text: str) -> tuple[str, list[str]]:
    expanded = text
    replacements: list[str] = []
    for abbreviation, full_text in sorted(
        ABBREVIATIONS.items(),
        key=lambda item: len(item[0]),
        reverse=True,
    ):
        pattern = re.compile(
            rf"(?<![\w/-]){re.escape(abbreviation)}(?![\w-])",
            re.IGNORECASE,
        )
        if pattern.search(expanded):
            expanded = pattern.sub(full_text, expanded)
            replacements.append(f"{abbreviation}->{full_text}")
    return expanded, replacements


def _normalize_for_matching(text: str) -> str:
    normalized = unicodedata.normalize("NFKD", text)
    without_marks = "".join(
        char for char in normalized if not unicodedata.combining(char)
    )
    without_marks = without_marks.replace("đ", "d").replace("Đ", "D")
    return " ".join(without_marks.casefold().split())


def _extract_refs(normalized_text: str) -> dict[str, str]:
    refs: dict[str, str] = {}
    if article_match := ARTICLE_PATTERN.search(normalized_text):
        refs["article"] = f"Dieu {article_match.group(1)}"
    if clause_match := CLAUSE_PATTERN.search(normalized_text):
        refs["clause"] = f"Khoan {clause_match.group(1)}"
    if point_match := POINT_PATTERN.search(normalized_text):
        refs["point"] = f"Diem {point_match.group(1)}"
    return refs