"""
TextNormalizer — lớp chuẩn hóa truy vấn tiếng Việt (khối "Text Normalization"
trong pipeline của bài báo, Fig. 2).

Người dùng thật gõ thiếu dấu, viết tắt và sai hoa thường rất thường xuyên
("quân binh thanh", "tp.hcm", "q.1"...). Lớp này chạy TRƯỚC NER/Intent để các
tầng sau nhận được văn bản sạch:

  1. Chuẩn hóa Unicode (NFC) + gom khoảng trắng thừa.
  2. Bung viết tắt phổ biến:  tp.hcm → Thành phố Hồ Chí Minh,
     q.1/q1 → Quận 1, p.5 → Phường 5, h. → Huyện, tx. → Thị xã ...
  3. Phục hồi dấu cho ĐỊA DANH bằng gazetteer dựng từ datasets/Wards.csv và
     Provinces.csv:  "quan binh thanh" → "Quận Bình Thạnh",
     "phuong linh tay" → "Phường Linh Tây". So khớp không phân biệt
     dấu/hoa-thường, ưu tiên cụm dài nhất, chỉ thay khi đúng ranh giới từ.

Mọi phép biến đổi được ghi lại trong `NormalizedQuery.changes` để có thể
hiển thị/giải thích (provenance).
"""

from __future__ import annotations

import csv
import logging
import re
import unicodedata
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

_REPO_ROOT = Path(__file__).resolve().parents[2]
_DATASETS = _REPO_ROOT / "datasets"

_TYPE_TITLE = {
    "phường": "Phường", "xã": "Xã", "thị trấn": "Thị trấn", "thị xã": "Thị xã",
    "quận": "Quận", "huyện": "Huyện", "thành phố": "Thành phố",
}


def fold(s: str) -> str:
    """Bỏ dấu + thường hóa, ánh xạ 1:1 từng ký tự (giữ nguyên offset)."""
    s = (s or "").replace("đ", "d").replace("Đ", "D")
    nfd = unicodedata.normalize("NFD", s)
    return "".join(c for c in nfd if unicodedata.category(c) != "Mn").lower()


def _cap_type_phrase(phrase: str) -> str:
    """'quận bình thạnh' -> 'Quận Bình Thạnh' (viết hoa từ loại + tên riêng)."""
    phrase = (phrase or "").strip()
    low = phrase.lower()
    for t in sorted(_TYPE_TITLE, key=len, reverse=True):
        if low.startswith(t + " ") or low == t:
            rest = phrase[len(t):].strip()
            rest_cap = " ".join(w[:1].upper() + w[1:] for w in rest.split())
            return f"{_TYPE_TITLE[t]} {rest_cap}".strip()
    return phrase


@dataclass
class NormalizedQuery:
    original: str
    text: str
    changes: List[Tuple[str, str, str]] = field(default_factory=list)  # (from, to, kind)

    @property
    def changed(self) -> bool:
        return self.text != self.original


class TextNormalizer:
    """Singleton qua get_text_normalizer(). Gazetteer build một lần (~5k mục)."""

    # (pattern, replacement, kind) — chạy theo thứ tự, case-insensitive
    _ABBREVIATIONS: List[Tuple[str, str, str]] = [
        (r"\btp\.?\s*hcm\b", "Thành phố Hồ Chí Minh", "abbrev"),
        (r"\btphcm\b", "Thành phố Hồ Chí Minh", "abbrev"),
        (r"\bhcm\b", "Hồ Chí Minh", "abbrev"),
        (r"\btp\.\s*(?=\w)", "Thành phố ", "abbrev"),
        (r"\btx\.\s*(?=\w)", "Thị xã ", "abbrev"),
        (r"\btt\.\s*(?=\w)", "Thị trấn ", "abbrev"),
        (r"\bq\.?\s*(\d{1,2})\b", r"Quận \1", "abbrev"),
        (r"\bp\.?\s*(\d{1,2})\b", r"Phường \1", "abbrev"),
        (r"\bq\.\s*(?=\w)", "Quận ", "abbrev"),
        (r"\bp\.\s*(?=\w)", "Phường ", "abbrev"),
        (r"\bh\.\s*(?=\w)", "Huyện ", "abbrev"),
    ]

    def __init__(self) -> None:
        # folded phrase -> canonical phrase; sắp theo độ dài giảm dần khi quét
        self._gazetteer: Dict[str, str] = {}
        self._build_gazetteer()
        self._entries = sorted(self._gazetteer.items(), key=lambda kv: -len(kv[0]))

    # ------------------------------------------------------------- gazetteer
    def _add(self, canonical: str, min_len: int = 0) -> None:
        canonical = (canonical or "").strip()
        f = fold(canonical)
        if len(f) >= max(min_len, 3) and f not in self._gazetteer:
            self._gazetteer[f] = canonical

    def _build_gazetteer(self) -> None:
        wards_csv = _DATASETS / "Wards.csv"
        if wards_csv.exists():
            try:
                with open(wards_csv, encoding="utf-8") as fh:
                    for row in csv.DictReader(fh):
                        # quận/huyện (kèm từ loại + dạng trần đủ dài)
                        dc = (row.get("district_city") or "").strip()
                        if dc:
                            self._add(_cap_type_phrase(dc))
                            bare = re.sub(
                                r"^(quận|huyện|thị xã|thành phố)\s+", "", dc, flags=re.IGNORECASE
                            ).strip()
                            if bare and not bare.isdigit():
                                self._add(" ".join(
                                    w[:1].upper() + w[1:] for w in bare.split()), min_len=6)
                        # phường/xã cũ và mới (kèm từ loại)
                        for tcol, ncol in (("old_type", "old_name"), ("new_type", "new_name")):
                            t = (row.get(tcol) or "").strip()
                            n = (row.get(ncol) or "").strip()
                            if t and n:
                                self._add(_cap_type_phrase(f"{t} {n}"))
            except Exception as e:  # pragma: no cover
                logger.warning(f"TextNormalizer: không đọc được Wards.csv: {e}")

        prov_csv = _DATASETS / "Provinces.csv"
        if prov_csv.exists():
            try:
                with open(prov_csv, encoding="utf-8") as fh:
                    for row in csv.DictReader(fh):
                        for key in ("old_province_name", "new_province_name"):
                            name = (row.get(key) or "").strip()
                            if name:
                                self._add(name, min_len=5)
                                self._add(f"tỉnh {name}" if name not in
                                          ("Hồ Chí Minh", "Cần Thơ") else f"Thành phố {name}")
            except Exception as e:  # pragma: no cover
                logger.warning(f"TextNormalizer: không đọc được Provinces.csv: {e}")

        logger.info(f"TextNormalizer: gazetteer {len(self._gazetteer)} mục")

    # ------------------------------------------------------------- pipeline
    def normalize(self, text: str) -> NormalizedQuery:
        original = text or ""
        changes: List[Tuple[str, str, str]] = []

        # 1) Unicode NFC + gom khoảng trắng
        out = unicodedata.normalize("NFC", original)
        out = re.sub(r"\s+", " ", out).strip()

        # 1.5) Câu VIẾT HOA TOÀN BỘ là out-of-distribution với NER/Intent
        # (train data là chữ thường/title-case) -> hạ về chữ thường; bước
        # phục hồi gazetteer phía dưới sẽ dựng lại đúng tên riêng có dấu.
        letters = [c for c in out if c.isalpha()]
        if len(letters) >= 8 and sum(c.isupper() for c in letters) / len(letters) > 0.7:
            changes.append((out, out.lower(), "lowercase_allcaps"))
            out = out.lower()

        # 2) Bung viết tắt
        for pattern, repl, kind in self._ABBREVIATIONS:
            def _sub(m: "re.Match[str]") -> str:
                new = m.expand(repl) if "\\" in repl or "$" in repl else repl
                changes.append((m.group(0), new, kind))
                return new
            out = re.sub(pattern, _sub, out, flags=re.IGNORECASE)

        # 3) Phục hồi dấu địa danh theo gazetteer
        out = self._restore_diacritics(out, changes)

        return NormalizedQuery(original=original, text=out, changes=changes)

    def _restore_diacritics(self, text: str, changes: List[Tuple[str, str, str]]) -> str:
        folded = fold(text)  # 1:1 — offset trùng với text
        n = len(folded)
        taken: List[Tuple[int, int]] = []
        repls: List[Tuple[int, int, str]] = []

        def _boundary_ok(s: int, e: int) -> bool:
            before = folded[s - 1] if s > 0 else " "
            after = folded[e] if e < n else " "
            return not before.isalnum() and not after.isalnum()

        def _overlaps(s: int, e: int) -> bool:
            return any(not (e <= ts or s >= te) for ts, te in taken)

        for f_phrase, canonical in self._entries:  # dài trước, ngắn sau
            start = 0
            while True:
                idx = folded.find(f_phrase, start)
                if idx == -1:
                    break
                end = idx + len(f_phrase)
                if _boundary_ok(idx, end) and not _overlaps(idx, end):
                    span = text[idx:end]
                    if span != canonical:  # chỉ thay khi thật sự khác (thiếu dấu/sai hoa thường)
                        repls.append((idx, end, canonical))
                        changes.append((span, canonical, "diacritics"))
                    taken.append((idx, end))
                start = end

        for s, e, canonical in sorted(repls, key=lambda r: -r[0]):  # thay từ phải sang trái
            text = text[:s] + canonical + text[e:]
        return text

    # ------------------------------------------------------------- tiện ích
    def find_district(self, text: str) -> Optional[str]:
        """Tên quận/huyện chuẩn xuất hiện trong câu (sau khi đã normalize)."""
        folded = fold(text)
        best: Optional[str] = None
        for f_phrase, canonical in self._entries:
            if fold(canonical).startswith(("quan ", "huyen ", "thi xa ", "thanh pho ")):
                if f_phrase in folded and (best is None or len(f_phrase) > len(fold(best))):
                    best = canonical
        return best

    def fold(self, s: str) -> str:
        return fold(s)


# Singleton
_normalizer: Optional[TextNormalizer] = None


def get_text_normalizer() -> TextNormalizer:
    global _normalizer
    if _normalizer is None:
        _normalizer = TextNormalizer()
    return _normalizer