import gradio as gr
import json
import os
import re
import gc
import ipaddress
import unicodedata
from datetime import datetime
from typing import Dict, List, Optional, Tuple, Callable, Pattern
from collections import defaultdict
from dataclasses import dataclass
from urllib.parse import urlparse

import torch
from gliner2 import GLiNER2
from gliner2.processor import WhitespaceTokenSplitter
from transformers import pipeline
from huggingface_hub import snapshot_download

# ==========================================
# 1. Custom CSS
# ==========================================
custom_css = """
.pii-entity {
    font-weight: bold;
    background-color: rgba(250, 204, 21, 0.3);
    border-bottom: 2px solid #eab308;
    border-radius: 4px;
    padding: 2px 4px;
    position: relative;
    cursor: help;
    transition: background-color 0.2s;
}
.pii-entity:hover {
    background-color: rgba(250, 204, 21, 0.6);
}
.pii-entity .tooltip-label {
    visibility: hidden;
    background-color: #1f2937;
    color: #f9fafb;
    text-align: center;
    border-radius: 6px;
    padding: 4px 8px;
    position: absolute;
    z-index: 50;
    bottom: 125%;
    left: 50%;
    transform: translateX(-50%);
    opacity: 0;
    transition: opacity 0.2s, bottom 0.2s;
    font-size: 0.75rem;
    font-family: monospace;
    font-weight: normal;
    white-space: nowrap;
    box-shadow: 0 4px 6px -1px rgba(0,0,0,0.1), 0 2px 4px -1px rgba(0,0,0,0.06);
}
.pii-entity .tooltip-label::after {
    content: "";
    position: absolute;
    top: 100%;
    left: 50%;
    margin-left: -5px;
    border-width: 5px;
    border-style: solid;
    border-color: #1f2937 transparent transparent transparent;
}
.pii-entity:hover .tooltip-label {
    visibility: visible;
    opacity: 1;
    bottom: 135%;
}
.text-container {
    font-size: 1rem;
    line-height: 1.8;
    padding: 1.5rem;
    background: #ffffff;
    border: 1px solid #e5e7eb;
    border-radius: 8px;
    min-height: 120px;
}
.dark .text-container {
    background: #1f2937;
    border-color: #374151;
    color: #e5e7eb;
}
.model-header-gliner {
    background: linear-gradient(135deg, #3b82f6, #1d4ed8);
    color: white;
    padding: 8px 16px;
    border-radius: 8px 8px 0 0;
    font-weight: bold;
    font-size: 0.9rem;
    margin-bottom: 0;
}
.model-header-xlmbert {
    background: linear-gradient(135deg, #8b5cf6, #5b21b6);
    color: white;
    padding: 8px 16px;
    border-radius: 8px 8px 0 0;
    font-weight: bold;
    font-size: 0.9rem;
    margin-bottom: 0;
}
.model-header-mmbert {
    background: linear-gradient(135deg, #f97316, #c2410c);
    color: white;
    padding: 8px 16px;
    border-radius: 8px 8px 0 0;
    font-weight: bold;
    font-size: 0.9rem;
    margin-bottom: 0;
}
.batch-sample {
    border: 1px solid #e5e7eb;
    border-radius: 10px;
    margin-bottom: 24px;
    overflow: hidden;
    box-shadow: 0 1px 3px rgba(0,0,0,0.06);
}
.dark .batch-sample {
    border-color: #374151;
}
.batch-sample-header {
    background: linear-gradient(135deg, #0ea5e9, #0284c7);
    color: white;
    padding: 10px 16px;
    font-weight: bold;
    font-size: 0.95rem;
    display: flex;
    justify-content: space-between;
    align-items: center;
}
.batch-sample-input {
    padding: 12px 16px;
    background: #f8fafc;
    border-bottom: 1px solid #e5e7eb;
    font-size: 0.85rem;
    color: #475569;
    white-space: pre-wrap;
    max-height: 120px;
    overflow-y: auto;
}
.dark .batch-sample-input {
    background: #111827;
    border-color: #374151;
    color: #94a3b8;
}
.batch-models-row {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 0;
}
.batch-model-col {
    padding: 12px 16px;
}
.batch-model-col:first-child {
    border-right: 1px solid #e5e7eb;
}
.dark .batch-model-col:first-child {
    border-color: #374151;
}
.batch-model-label {
    font-size: 0.75rem;
    font-weight: 600;
    text-transform: uppercase;
    letter-spacing: 0.05em;
    margin-bottom: 8px;
    padding: 3px 8px;
    border-radius: 4px;
    display: inline-block;
}
.batch-model-label.gliner {
    background: #dbeafe;
    color: #1d4ed8;
}
.batch-model-label.bert {
    background: #ede9fe;
    color: #5b21b6;
}
.batch-model-label.mmbert {
    background: #ffedd5;
    color: #c2410c;
}
.dark .batch-model-label.gliner {
    background: #1e3a5f;
    color: #93c5fd;
}
.dark .batch-model-label.bert {
    background: #2e1065;
    color: #c4b5fd;
}
.dark .batch-model-label.mmbert {
    background: #431407;
    color: #fdba74;
}
.batch-entity-tags {
    display: flex;
    flex-wrap: wrap;
    gap: 4px;
    margin-top: 6px;
}
.batch-entity-tag {
    font-size: 0.72rem;
    padding: 2px 6px;
    border-radius: 4px;
    background: #fef3c7;
    color: #92400e;
    font-family: monospace;
    border: 1px solid #fde68a;
}
.dark .batch-entity-tag {
    background: #422006;
    color: #fde68a;
    border-color: #78350f;
}
.batch-summary {
    background: linear-gradient(135deg, #f0fdf4, #dcfce7);
    border: 1px solid #bbf7d0;
    border-radius: 10px;
    padding: 16px 20px;
    margin-bottom: 20px;
    font-size: 0.9rem;
}
.dark .batch-summary {
    background: linear-gradient(135deg, #052e16, #064e3b);
    border-color: #166534;
    color: #bbf7d0;
}
.batch-progress {
    text-align: center;
    padding: 40px;
    color: #6b7280;
    font-size: 1.1rem;
}
"""

# ==========================================
# 2. Cấu hình Label & Tokenizer (Cho GLiNER)
# ==========================================
_IMPROVED_PATTERN = re.compile(
    r"""(?:https?://[^\s]+|www\.[^\s]+)|[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}|@[a-z0-9_]+|(?:[0-9a-f]{1,4}:){7}[0-9a-f]{1,4}|(?:[0-9a-f]{1,4}:){1,7}:|:(?::[0-9a-f]{1,4}){1,7}|[a-z0-9_-]{10,}\.[a-z0-9_-]{10,}\.[a-z0-9_-]{10,}|\w+(?:[-_]\w+)*|\S""",
    re.VERBOSE | re.IGNORECASE,
)
WhitespaceTokenSplitter._PATTERN = _IMPROVED_PATTERN

LABELS = {
    "PREFIX": "Personal titles and prefixes",
    "PERSONAL_NAME": "Full, given, and family names",
    "SEX": "Biological sex and gender identities",
    "AGE": "Age values",
    "DOB": "Dates of birth",
    "PHONE_NUMBER": "Telephone and mobile numbers",
    "EMAIL": "Email addresses",
    "LOCATION": "Geographic regions and cities",
    "STREET_ADDRESS": "Specific local and street addresses",
    "ZIPCODE": "Postal and ZIP codes",
    "GPS_COORDINATE": "GPS coordinates and latitude/longitude",
    "USERNAME": "Usernames and digital aliases",
    "PASSWORD": "Passwords and login credentials",
    "PIN": "Personal Identification Numbers (PINs)",
    "URL": "URLs and web addresses",
    "IP_ADDRESS": "IP addresses",
    "ACCOUNT_NUMBER": "Bank account numbers",
    "AMOUNT": "Monetary amounts and values",
    "CREDIT_CARD_ISSUER": "Credit card issuers and brands",
    "CREDIT_CARD_NUMBER": "Credit card numbers",
    "CREDIT_CARD_CVV": "Credit card verification values (CVV/CVC)",
    "IBAN": "International Bank Account Numbers (IBAN)",
    "BIC_SWIFT": "BIC and SWIFT codes",
    "CRYPTO_ADDRESS": "Cryptocurrency addresses",
    "OCCUPATION": "Occupations, job titles, and industries",
    "SSN_CCCD": "Social Security and national identification numbers",
    "PASSPORT_NUM": "Passport numbers",
    "DRIVER_LICENSE": "Driver's license numbers",
    "TAX_ID": "Tax identification numbers",
    "DATE": "General calendar dates",
    "TIME": "Specific times of day",
    "MARITAL_STATUS": "Marital and legal relationship status",
    "RELIGION": "Religious affiliations and beliefs",
    "ETHNICITY": "Ethnicities and cultural backgrounds",
    "TRADE_UNION_INFO": "Trade union membership and affiliations",
    "NATIONALITY": "Nationalities and citizenships",
    "HEALTH_INSURANCE": "Health insurance numbers and identifiers",
    "HEALTH_STATUS": "Health status and medical conditions"
}

# ==========================================
# 2b. Preprocessing — clean_text_with_mapping
# ==========================================
_STRIP_CHARS = frozenset(["\u00ad", "\u200b", "\u200c", "\u200d", "\u2060", "\ufeff", "\ufffd"])
_C1_RE = re.compile(r"[\x80-\x9f]")
_C0_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
_MULTI_SPACE_RE = re.compile(r"[^\S\n]+")


def clean_text_with_mapping(text: str) -> tuple:
    if not text:
        return text, []
    text = unicodedata.normalize("NFC", text)
    cleaned_chars = []
    char_map = []
    orig_pos = 0
    for ch in text:
        if ch in _STRIP_CHARS:
            orig_pos += 1
            continue
        if _C0_RE.match(ch):
            orig_pos += 1
            continue
        if _C1_RE.match(ch):
            cleaned_chars.append(" ")
            char_map.append(orig_pos)
            orig_pos += 1
            continue
        cleaned_chars.append(ch)
        char_map.append(orig_pos)
        orig_pos += 1
    cleaned = "".join(cleaned_chars)
    cleaned = _MULTI_SPACE_RE.sub(" ", cleaned).strip()
    final_map = []
    i = 0
    for ch in cleaned:
        while i < len(char_map) and (text[char_map[i]] != ch and not (ch == " " and text[char_map[i]].isspace())):
            i += 1
        final_map.append(char_map[i] if i < len(char_map) else len(text))
        i += 1
    final_map.append(len(text))
    return cleaned, final_map


# ==========================================
# 2c. Rule-based detector
# ==========================================
@dataclass(frozen=True)
class RegexRule:
    entity_type: str
    pattern: Pattern[str]
    validator: Callable[[str], bool]
    normalizer: Callable[[str], str] | None = None


MONTH_PATTERN = (
    r"(?:Jan(?:uary)?|Feb(?:ruary)?|Mar(?:ch)?|Apr(?:il)?|May|"
    r"Jun(?:e)?|Jul(?:y)?|Aug(?:ust)?|Sep(?:t(?:ember)?)?|"
    r"Oct(?:ober)?|Nov(?:ember)?|Dec(?:ember)?)"
)
MONTH_PATTERN_VI = r"(?:tháng\s*(?:1|2|3|4|5|6|7|8|9|10|11|12|một|hai|ba|tư|năm|sáu|bảy|tám|chín|mười|mười\s*một|mười\s*hai))"
MONTH_PATTERN_DE = r"(?:Januar|Februar|März|April|Mai|Juni|Juli|August|September|Oktober|November|Dezember)"
DATE_TOKEN_PATTERN = (
    rf"(?:\d{{4}}[-/]\d{{1,2}}[-/]\d{{1,2}}|"
    rf"\d{{1,2}}[-/]\d{{1,2}}[-/]\d{{2,4}}|"
    rf"\d{{1,2}}\.\d{{1,2}}\.\d{{2,4}}|"
    rf"{MONTH_PATTERN}\s+\d{{1,2}}(?:st|nd|rd|th)?(?:,\s*\d{{2,4}})?|"
    rf"\d{{1,2}}(?:st|nd|rd|th)?\s+(?:of\s+)?{MONTH_PATTERN}(?:,\s*\d{{2,4}})?|"
    rf"(?:ngày\s+)?\d{{1,2}}\s+{MONTH_PATTERN_VI}(?:\s+(?:năm\s+)?\d{{2,4}})?|"
    rf"\d{{1,2}}\.\s*{MONTH_PATTERN_DE}(?:\s+\d{{2,4}})?)"
)


def normalize_space(value: str) -> str:
    return re.sub(r"\s+", " ", value).strip()

def strip_number_separators(value: str) -> str:
    return re.sub(r"[\s-]", "", value).strip()

def normalize_phone(value: str) -> str:
    compact = re.sub(r"[().\s-]", "", value)
    if value.strip().startswith("+"):
        return "+" + compact.lstrip("+")
    return compact

def normalize_iban(value: str) -> str:
    return re.sub(r"[^A-Z0-9]", "", value.upper())


def is_valid_email(value: str) -> bool:
    if ".." in value:
        return False
    local, _, domain = value.rpartition("@")
    if not local or not domain or "." not in domain:
        return False
    return not domain.startswith(".") and not domain.endswith(".")

def is_valid_url(value: str) -> bool:
    candidate = value
    if candidate.lower().startswith("www."):
        candidate = "https://" + candidate
    parsed = urlparse(candidate)
    return bool(parsed.scheme and parsed.netloc and "." in parsed.netloc)

def is_valid_phone(value: str) -> bool:
    digits = re.sub(r"\D", "", value)
    has_separator = any(ch in value for ch in " +-().")
    groups = re.findall(r"\d+", value)
    if not (8 <= len(digits) <= 15 and has_separator):
        return False
    if any(len(group) > 4 for group in groups):
        return False
    if groups and len(groups[-1]) < 3:
        return False
    if not value.strip().startswith("+") and len(digits) > 11:
        return False
    if re.fullmatch(r"\d{4}-\d{1,2}-\d{1,2}(?:\s+\d{1,2})?", value):
        return False
    return True

def is_valid_ipv4(value: str) -> bool:
    try:
        return isinstance(ipaddress.ip_address(value), ipaddress.IPv4Address)
    except ValueError:
        return False

def is_valid_ipv6(value: str) -> bool:
    try:
        return isinstance(ipaddress.ip_address(value), ipaddress.IPv6Address)
    except ValueError:
        return False

def parse_date_text(value: str) -> bool:
    candidate = normalize_space(value.replace(" of ", " "))
    formats = [
        "%Y-%m-%d", "%Y/%m/%d", "%d-%m-%Y", "%d/%m/%Y", "%m-%d-%Y", "%m/%d/%Y",
        "%d-%m-%y", "%d/%m/%y", "%m-%d-%y", "%m/%d/%y",
        "%B %d, %Y", "%B %d, %y", "%B %d %Y", "%B %d",
        "%b %d, %Y", "%b %d, %y", "%b %d %Y", "%b %d",
        "%d %B %Y", "%d %B", "%d %b %Y", "%d %b",
    ]
    candidate = re.sub(r"(\d)(st|nd|rd|th)\b", r"\1", candidate, flags=re.IGNORECASE)
    for fmt in formats:
        try:
            datetime.strptime(candidate, fmt)
            return True
        except ValueError:
            continue
    return False

def is_valid_time(value: str) -> bool:
    candidate = normalize_space(value).upper().replace(".", "")
    for fmt in ["%H:%M", "%I:%M %p", "%I %p"]:
        try:
            datetime.strptime(candidate, fmt)
            return True
        except ValueError:
            continue
    return False

def luhn_check(value: str) -> bool:
    digits = re.sub(r"\D", "", value)
    if not 13 <= len(digits) <= 19:
        return False
    total = 0
    for index, digit in enumerate(digits[::-1]):
        number = int(digit)
        if index % 2 == 1:
            number *= 2
            if number > 9:
                number -= 9
        total += number
    return total % 10 == 0

IBAN_LENGTHS = {
    "AD":24,"AE":23,"AL":28,"AT":20,"AZ":28,"BA":20,"BE":16,"BG":22,"BH":22,"BR":29,
    "BY":28,"CH":21,"CR":22,"CY":28,"CZ":24,"DE":22,"DK":18,"DO":28,"EE":20,"EG":29,
    "ES":24,"FI":18,"FO":18,"FR":27,"GB":22,"GE":22,"GI":23,"GL":18,"GR":27,"GT":28,
    "HR":21,"HU":28,"IE":22,"IL":23,"IQ":23,"IS":26,"IT":27,"JO":30,"KW":30,"KZ":20,
    "LB":28,"LC":32,"LI":21,"LT":20,"LU":20,"LV":21,"MC":27,"MD":24,"ME":22,"MK":19,
    "MR":27,"MT":31,"MU":30,"NL":18,"NO":15,"PK":24,"PL":28,"PS":29,"PT":25,"QA":29,
    "RO":24,"RS":22,"SA":24,"SC":31,"SE":24,"SI":19,"SK":24,"SM":27,"ST":25,"SV":28,
    "TL":23,"TN":24,"TR":26,"UA":29,"VA":22,"VG":24,"XK":20,
}

def iban_checksum(value: str) -> bool:
    compact = normalize_iban(value)
    if not 15 <= len(compact) <= 34:
        return False
    expected = IBAN_LENGTHS.get(compact[:2])
    if expected and len(compact) != expected:
        return False
    rearranged = compact[4:] + compact[:4]
    numeric = "".join(str(int(ch, 36)) for ch in rearranged)
    return int(numeric) % 97 == 1

def is_valid_zipcode(value: str) -> bool:
    c = value.strip()
    return bool(
        re.fullmatch(r"\d{5}(?:-\d{4})?", c)
        or re.fullmatch(r"[A-Z]{1,2}\d[A-Z0-9]?\s*\d[A-Z]{2}", c, re.IGNORECASE)
        or re.fullmatch(r"\d{5}", c)
        or re.fullmatch(r"\d{6}", c)
        or re.fullmatch(r"[A-Z]\d[A-Z]\s*\d[A-Z]\d", c, re.IGNORECASE)
        or re.fullmatch(r"\d{3}-\d{4}", c)
        or re.fullmatch(r"\d{4}", c)
    )

def is_valid_amount(value: str) -> bool:
    digits = re.sub(r"[^\d.]", "", value)
    if not digits or digits.count(".") > 1:
        return False
    try:
        return float(digits) > 0
    except ValueError:
        return False

def is_valid_crypto_address(value: str) -> bool:
    c = value.strip()
    return bool(
        re.fullmatch(r"[13][A-HJ-NP-Za-km-z1-9]{25,34}", c)
        or re.fullmatch(r"bc1[a-z0-9]{25,89}", c)
        or re.fullmatch(r"0x[0-9a-fA-F]{40}", c)
        or re.fullmatch(r"[LM3][A-HJ-NP-Za-km-z1-9]{26,33}", c)
        or re.fullmatch(r"r[0-9A-HJ-NP-Za-km-z]{24,34}", c)
    )

def is_valid_age(value: str) -> bool:
    digits = re.sub(r"\D", "", value)
    return bool(digits) and 0 <= int(digits) <= 150

def is_valid_ssn(value: str) -> bool:
    digits = re.sub(r"\D", "", value)
    if len(digits) != 9:
        return False
    area, group, serial = digits[:3], digits[3:5], digits[5:]
    if area in {"000", "666"} or area.startswith("9"):
        return False
    return group != "00" and serial != "0000"

def is_valid_cccd(value: str) -> bool:
    digits = re.sub(r"\D", "", value)
    return len(digits) in {9, 12}


RULE_PATTERNS: Dict[str, RegexRule] = {
    "EMAIL": RegexRule(
        entity_type="EMAIL",
        pattern=re.compile(r"(?<![\w.+-])[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,63}(?![\w-])", re.IGNORECASE),
        validator=is_valid_email,
    ),
    "URL": RegexRule(
        entity_type="URL",
        pattern=re.compile(r"(?:(?:https?://)|(?:www\.))[^\s<>()]+", re.IGNORECASE),
        validator=is_valid_url,
    ),
    "PHONE_NUMBER": RegexRule(
        entity_type="PHONE_NUMBER",
        pattern=re.compile(r"(?<!\w)(?:\+?\d{1,3}[\s().-]*)?(?:\(?\d{2,4}\)?[\s().-]*){1,3}\d{3,4}(?!\w)"),
        validator=is_valid_phone,
        normalizer=normalize_phone,
    ),
    "IP_ADDRESS_V4": RegexRule(
        entity_type="IP_ADDRESS",
        pattern=re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"),
        validator=is_valid_ipv4,
    ),
    "IP_ADDRESS_V6": RegexRule(
        entity_type="IP_ADDRESS",
        pattern=re.compile(r"(?<![:\w])(?:[A-F0-9]{1,4}:){2,7}[A-F0-9]{1,4}(?![:\w])", re.IGNORECASE),
        validator=is_valid_ipv6,
    ),
    "DATE": RegexRule(
        entity_type="DATE",
        pattern=re.compile(rf"\b{DATE_TOKEN_PATTERN}\b", re.IGNORECASE),
        validator=parse_date_text,
    ),
    "TIME": RegexRule(
        entity_type="TIME",
        pattern=re.compile(
            r"\b(?:(?:[01]?\d|2[0-3]):[0-5]\d(?::[0-5]\d)?(?:\s?(?:AM|PM))?|(?:1[0-2]|0?[1-9])\s?(?:AM|PM))\b",
            re.IGNORECASE,
        ),
        validator=is_valid_time,
    ),
    "CREDIT_CARD_NUMBER": RegexRule(
        entity_type="CREDIT_CARD_NUMBER",
        pattern=re.compile(r"(?<!\d)(?:\d[ -]?){13,19}(?!\d)"),
        validator=luhn_check,
        normalizer=strip_number_separators,
    ),
    "IBAN": RegexRule(
        entity_type="IBAN",
        pattern=re.compile(r"\b[A-Z]{2}\d{2}(?:[ -]?[A-Z0-9]){11,30}\b", re.IGNORECASE),
        validator=iban_checksum,
        normalizer=normalize_iban,
    ),
    "SSN_CCCD_US": RegexRule(
        entity_type="SSN_CCCD",
        pattern=re.compile(r"\b(?!000|666|9\d\d)\d{3}-\d{2}-\d{4}\b"),
        validator=is_valid_ssn,
    ),
    "SSN_CCCD_VN": RegexRule(
        entity_type="SSN_CCCD",
        pattern=re.compile(r"\b0\d{2}[0-3]\d{8}\b"),
        validator=is_valid_cccd,
    ),
    "ZIPCODE": RegexRule(
        entity_type="ZIPCODE",
        pattern=re.compile(
            r"(?<=[\s,])\d{5}(?:-\d{4})?(?=[\s,.\n]|$)"
            r"|(?<=[\s,])[A-Z]{1,2}\d[A-Z0-9]?\s*\d[A-Z]{2}(?=[\s,.\n]|$)"
            r"|(?<=[\s,])[A-Z]\d[A-Z]\s*\d[A-Z]\d(?=[\s,.\n]|$)"
            r"|(?<=[\s,])\d{3}-\d{4}(?=[\s,.\n]|$)",
            re.IGNORECASE,
        ),
        validator=is_valid_zipcode,
    ),
    "AMOUNT": RegexRule(
        entity_type="AMOUNT",
        pattern=re.compile(
            r"(?<!\w)(?:[$€£¥₫₩₹])\s*\d[\d,. ]*\d(?:\s*(?:USD|EUR|GBP|VND|JPY|KRW|INR|AUD|CAD|CHF|CNY|SGD|THB|đồng|đ|dong))?"
            r"|(?<!\w)\d[\d,. ]*\d\s*(?:USD|EUR|GBP|VND|JPY|KRW|INR|AUD|CAD|CHF|CNY|SGD|THB|đồng|đ|dong)\b",
            re.IGNORECASE,
        ),
        validator=is_valid_amount,
    ),
    "CRYPTO_ADDRESS": RegexRule(
        entity_type="CRYPTO_ADDRESS",
        pattern=re.compile(
            r"\b(?:0x[0-9a-fA-F]{40}|[13][A-HJ-NP-Za-km-z1-9]{25,34}|bc1[a-z0-9]{25,89}"
            r"|[LM3][A-HJ-NP-Za-km-z1-9]{26,33}|r[0-9A-HJ-NP-Za-km-z]{24,34})\b"
        ),
        validator=is_valid_crypto_address,
    ),
    "AGE": RegexRule(
        entity_type="AGE",
        pattern=re.compile(
            r"(?i)(?:"
            r"(?:age(?:d)?|tuổi|năm tuổi|years?\s*old|ans|Jahre?\s*alt)\s*[:#]?\s*(\d{1,3})"
            r"|(\d{1,3})\s*(?:tuổi|năm tuổi|years?\s*old|ans|Jahre?\s*alt)"
            r")"
        ),
        validator=is_valid_age,
    ),
}

_RULE_PRIORITY: List[str] = [
    "CRYPTO_ADDRESS", "IBAN", "CREDIT_CARD_NUMBER",
    "SSN_CCCD_US", "SSN_CCCD_VN",
    "IP_ADDRESS_V6", "IP_ADDRESS_V4",
    "URL", "EMAIL", "DATE", "AMOUNT", "AGE",
]

RULE_BASED_LABELS: set = {rule.entity_type for rule in RULE_PATTERNS.values()}

_TRAILING_PUNCT = re.compile(r"[.,;:)\]}]+$")

def _trim_trailing_punct(value: str, end: int) -> tuple:
    trimmed = _TRAILING_PUNCT.sub("", value)
    return trimmed, end - (len(value) - len(trimmed))


def detect_by_rules(text: str, allowed_labels) -> List[Dict]:
    if isinstance(allowed_labels, dict):
        allowed_entity_types = set(allowed_labels.keys())
    else:
        allowed_entity_types = set(allowed_labels)
    results: List[Dict] = []
    for key in _RULE_PRIORITY:
        rule = RULE_PATTERNS.get(key)
        if rule is None or rule.entity_type not in allowed_entity_types:
            continue
        for m in rule.pattern.finditer(text):
            raw = m.group(0).strip()
            cleaned, end = _trim_trailing_punct(raw, m.end())
            if not cleaned:
                continue
            if not rule.validator(cleaned):
                continue
            results.append({
                "text": cleaned, "label": rule.entity_type,
                "start": m.start(), "end": end,
            })
    return results


# ==========================================
# 2d. Postprocessing — normalize & resolve conflicts
# ==========================================
def _find_all_occurrences(text: str, value: str) -> List[Tuple[int, int]]:
    spans, start = [], 0
    while True:
        idx = text.find(value, start)
        if idx == -1:
            break
        spans.append((idx, idx + len(value)))
        start = idx + max(1, len(value))
    return spans


def normalize_predictions(raw_preds, text: str) -> List[Dict]:
    if raw_preds is None:
        return []
    results: List[Dict] = []

    if isinstance(raw_preds, list):
        for item in raw_preds:
            if not isinstance(item, dict):
                continue
            label = item.get("label") or item.get("entity_group") or item.get("entity") or ""
            start = item.get("start", -1)
            end = item.get("end", -1)
            value = item.get("text") or item.get("value") or item.get("word") or ""
            if not label:
                continue
            score = float(item.get("score") or item.get("probability") or item.get("confidence") or 0.0)
            if start != -1 and end != -1 and 0 <= int(start) < int(end) <= len(text):
                results.append({"text": text[int(start):int(end)], "label": label,
                                "start": int(start), "end": int(end), "score": score})
            elif value:
                for s, e in _find_all_occurrences(text, str(value)):
                    results.append({"text": text[s:e], "label": label, "start": s, "end": e, "score": score})

    elif isinstance(raw_preds, dict):
        entities = raw_preds.get("entities", raw_preds)
        for label, values in entities.items():
            if not isinstance(values, list):
                continue
            for v in values:
                if isinstance(v, dict):
                    s, e = v.get("start"), v.get("end")
                    surface = v.get("text") or v.get("value") or ""
                    sc = float(v.get("score") or v.get("probability") or v.get("confidence") or 0.0)
                    if s is not None and e is not None and 0 <= int(s) < int(e) <= len(text):
                        results.append({"text": text[int(s):int(e)], "label": label,
                                        "start": int(s), "end": int(e), "score": sc})
                    elif surface:
                        for cs, ce in _find_all_occurrences(text, str(surface)):
                            results.append({"text": text[cs:ce], "label": label,
                                            "start": cs, "end": ce, "score": sc})
                else:
                    for cs, ce in _find_all_occurrences(text, str(v)):
                        results.append({"text": text[cs:ce], "label": label,
                                        "start": cs, "end": ce, "score": 0.0})

    best: Dict[tuple, Dict] = {}
    for r in results:
        key = (r["label"], r["start"], r["end"])
        if key not in best or r.get("score", 0.0) > best[key].get("score", 0.0):
            best[key] = r
    return sorted(best.values(), key=lambda x: x["start"])


_LABEL_PRIORITY: Dict[str, int] = {
    "CRYPTO_ADDRESS": 10, "IBAN": 10, "CREDIT_CARD_NUMBER": 10,
    "SSN_CCCD_US": 10, "SSN_CCCD_VN": 10,
    "IP_ADDRESS_V6": 10, "IP_ADDRESS_V4": 10,
    "URL": 10, "EMAIL": 10, "DATE": 10, "AMOUNT": 10, "AGE": 10,
}


def resolve_span_conflicts(
    entities: List[Dict],
    strategy: str = "score_first",
    rule_labels: Optional[set] = None,
) -> List[Dict]:
    if not entities:
        return []
    _rule_labels = rule_labels if rule_labels is not None else RULE_BASED_LABELS

    def _sort_key(e: Dict):
        is_rule = 1 if e["label"] in _rule_labels else 0
        score = 1.0 if is_rule else float(e.get("score") or 0.0)
        length = e["end"] - e["start"]
        prio = _LABEL_PRIORITY.get(e["label"], 0)
        if strategy == "score_first":
            return (is_rule, score, prio, length)
        elif strategy == "longest_first":
            return (is_rule, length, score, prio)
        else:
            return (is_rule, prio, score, length)

    ranked = sorted(entities, key=_sort_key, reverse=True)
    kept: List[Dict] = []
    occupied: List[Tuple[int, int]] = []

    def _overlaps(s: int, e: int) -> bool:
        for (ks, ke) in occupied:
            if s < ke and e > ks:
                return True
        return False

    for ent in ranked:
        s, e = ent["start"], ent["end"]
        if not _overlaps(s, e):
            kept.append(ent)
            occupied.append((s, e))

    return sorted(kept, key=lambda x: x["start"])


def apply_advanced_postprocessing(entities: List[Dict], original_text: str) -> List[Dict]:
    if not entities:
        return []
    sorted_ents = sorted(entities, key=lambda x: x['start'])
    final_entities = []
    for i, ent in enumerate(sorted_ents):
        label = ent["label"]
        start = ent["start"]
        end = ent["end"]
        if label in ["TIME", "DATE"]:
            prev_char = original_text[start-1] if start > 0 else ""
            next_char = original_text[end] if end < len(original_text) else ""
            if prev_char == "[" and next_char == "]":
                continue
        if label == "PREFIX":
            has_person_after = False
            if i + 1 < len(sorted_ents):
                next_ent = sorted_ents[i+1]
                if next_ent["label"] == "PERSONAL_NAME":
                    gap = next_ent["start"] - end
                    if 0 <= gap <= 3:
                        has_person_after = True
            if not has_person_after:
                continue
        final_entities.append(ent)
    return final_entities


# ==========================================
# 3. Model Wrappers
# ==========================================
class _BaseWrapper:
    def cleanup(self):
        if hasattr(self, "model"):
            del self.model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

    def _run_model(self, text: str, label_input):
        raw = None
        for method in ("predict_entities", "extract_entities"):
            fn = getattr(self.model, method, None)
            if fn is None:
                continue
            try:
                raw = fn(text, label_input, threshold=self.threshold)
            except TypeError:
                raw = fn(text, label_input)
            break
        return raw


class GLiNER2FinetunedWrapper(_BaseWrapper):
    def __init__(
        self,
        finetuned_repo: str = "quynong/gliner2-rosenxt",
        checkpoint: str = "best",
        use_label_description: bool = False,
        use_rules: bool = True,
        threshold: float = 0.3,
        MAX_SPAN: int = 100,
    ):
        self.use_label_description = use_label_description
        self.use_rules = use_rules
        self.threshold = threshold
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        print(f"Downloading {checkpoint!r} from {finetuned_repo}")
        local_path = snapshot_download(
            repo_id=finetuned_repo,
            allow_patterns=f"{checkpoint}/*"
        )
        checkpoint_path = os.path.join(local_path, checkpoint)
        print(f"Loading model from: {checkpoint_path}")
        self.model = GLiNER2.from_pretrained(checkpoint_path)
        self.model.config.max_width = MAX_SPAN
        self.model.max_width = MAX_SPAN
        self.model.span_rep.span_rep_layer.max_width = MAX_SPAN
        self.model.to(self.device).eval()

    def predict(self, text: str, labels, label_descriptions: Optional[Dict[str, str]] = None) -> List[Dict]:
        cleaned_text, char_map = clean_text_with_mapping(text)
        raw = self._run_model(
            cleaned_text,
            labels if not (self.use_label_description and label_descriptions)
            else {k: v for k, v in label_descriptions.items()
                  if k in (labels if isinstance(labels, set)
                            else set(labels.keys()) if isinstance(labels, dict)
                            else set(labels))}
        )
        model_preds = normalize_predictions(raw, cleaned_text)

        rule_preds: List[Dict] = []
        if self.use_rules:
            for e in detect_by_rules(cleaned_text, labels):
                rule_preds.append({**e, "score": 1.0})

        combined = rule_preds + model_preds
        resolved = resolve_span_conflicts(combined, strategy="rule_first")

        for pred in resolved:
            s, e = pred["start"], pred["end"]
            pred["start"] = char_map[s] if s < len(char_map) else len(text)
            pred["end"] = char_map[e] if e < len(char_map) else len(text)
            pred["text"] = text[pred["start"]:pred["end"]]

        return apply_advanced_postprocessing(resolved, text)


class mmBERTWrapper(_BaseWrapper):
    def __init__(
        self,
        hf_repo: str,
        subfolder: str = "",
        threshold: float = 0.5,
        use_rules: bool = True,
    ):
        self.threshold = threshold
        self.use_rules = use_rules

        if subfolder:
            print(f"Downloading subfolder '{subfolder}' from {hf_repo}")
            local_repo_path = snapshot_download(
                repo_id=hf_repo,
                allow_patterns=f"{subfolder}/*"
            )
            model_path = os.path.join(local_repo_path, subfolder)
        else:
            model_path = hf_repo

        print(f"Loading mmBERT pipeline from: {model_path}")
        self.pipe = pipeline(
            "ner",
            model=model_path,
            tokenizer=model_path,
            aggregation_strategy="simple",
        )

    def predict(self, text: str, labels) -> List[Dict]:
        cleaned_text, char_map = clean_text_with_mapping(text)
        raw = self.pipe(cleaned_text)
        model_preds = normalize_predictions(raw, cleaned_text)

        # Filter by threshold
        model_preds = [p for p in model_preds if p.get("score", 0.0) >= self.threshold]

        rule_preds: List[Dict] = []
        if self.use_rules:
            for e in detect_by_rules(cleaned_text, labels):
                rule_preds.append({**e, "score": 1.0})

        combined = rule_preds + model_preds
        resolved = resolve_span_conflicts(combined, strategy="rule_first")

        for pred in resolved:
            s, e = pred["start"], pred["end"]
            pred["start"] = char_map[s] if s < len(char_map) else len(text)
            pred["end"] = char_map[e] if e < len(char_map) else len(text)
            pred["text"] = text[pred["start"]:pred["end"]]

        return apply_advanced_postprocessing(resolved, text)

    def cleanup(self):
        if hasattr(self, "pipe"):
            del self.pipe
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()


# ==========================================
# 3b. Lazy model loaders (singleton)
# ==========================================
_MODEL_INSTANCES: Dict[str, object] = {}
_MODEL_FACTORIES = {
    "multi_1": lambda: GLiNER2FinetunedWrapper(
        finetuned_repo="quynong/gliner2-rosenxt-multi-v7",
        use_label_description=False,
        use_rules=False,
        checkpoint="checkpoint-epoch-3",
        threshold=0.7,
        MAX_SPAN=90,
    ),
    "XLMBERT": lambda: mmBERTWrapper(
        hf_repo="quynong/multilingualv1",
        subfolder="FacebookAI-xlm-roberta-base",
        threshold=0.5,
        use_rules=False,
    ),
    "mmBERT-multilingual": lambda: mmBERTWrapper(
        hf_repo="quynong/multilingual",
        subfolder="jhu-clsp-mmBERT-base",
        threshold=0.5,
        use_rules=False,
    ),
}

MODEL_DISPLAY_NAMES = {
    "multi_1": "GLiNER2 — multi-v7 (ep3)",
    "XLMBERT": "XLM-RoBERTa (BERT)",
    "mmBERT-multilingual": "mmBERT-multilingual",
}


def get_model(key: str):
    if key not in _MODEL_INSTANCES:
        print(f"🚀 Khởi tạo model: {key}")
        _MODEL_INSTANCES[key] = _MODEL_FACTORIES[key]()
        print(f"✅ Model '{key}' sẵn sàng.")
    return _MODEL_INSTANCES[key]


# ==========================================
# 4. Hàm tạo HTML
# ==========================================
def build_html_with_tooltips(text: str, spans: List[Dict]) -> str:
    spans = sorted(spans, key=lambda x: x["start"])
    html_content = ""
    last_idx = 0
    for span in spans:
        start, end, label = span["start"], span["end"], span["entity"]
        if start >= end or start < last_idx:
            continue
        html_content += text[last_idx:start].replace('\n', '<br>')
        pii_text = text[start:end].replace('\n', '<br>')
        html_content += (
            f'<span class="pii-entity">{pii_text}'
            f'<span class="tooltip-label">{label}</span></span>'
        )
        last_idx = end
    html_content += text[last_idx:].replace('\n', '<br>')
    return f'<div class="text-container">{html_content}</div>'


# ==========================================
# 5. Core: chạy 1 model, trả về (html, json)
# ==========================================
def run_single_model(text: str, model_key: str):
    try:
        wrapper = get_model(model_key)
        preds = wrapper.predict(text=text, labels=LABELS)
        spans_for_html = [{"entity": p["label"], "start": p["start"], "end": p["end"]} for p in preds]
        html = build_html_with_tooltips(text, spans_for_html)
        return html, preds
    except Exception as exc:
        import traceback
        tb = traceback.format_exc()
        return f"<div style='color:red;'>Lỗi model {model_key}: {exc}</div>", {"error": str(exc), "traceback": tb}


# ==========================================
# 6. Wrapper cho Gradio (3 model song song)
# ==========================================
def process_both(text: str):
    html1, json1 = run_single_model(text, "multi_1")
    html2, json2 = run_single_model(text, "XLMBERT")
    html3, json3 = run_single_model(text, "mmBERT-multilingual")
    return html1, json1, html2, json2, html3, json3


# ==========================================
# 6b. Batch processing — test file upload
# ==========================================
def _build_entity_tags_html(preds: list) -> str:
    if not preds or (isinstance(preds, dict) and "error" in preds):
        return '<span style="color:#ef4444;">Lỗi</span>'
    if not preds:
        return '<span style="color:#9ca3af;">Không phát hiện PII</span>'
    tags = []
    for p in preds:
        label = p.get("label", "?")
        text_val = p.get("text", "")
        display = text_val if len(text_val) <= 30 else text_val[:27] + "..."
        display = display.replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;")
        tags.append(f'<span class="batch-entity-tag" title="{display}">{label}: {display}</span>')
    return "".join(tags)


def process_test_file(file, model_choice: str):
    if file is None:
        return "<div class='batch-progress'>⬆️ Vui lòng upload file JSON.</div>", []

    try:
        file_path = file if isinstance(file, str) else file.name
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)
    except Exception as exc:
        return f"<div style='color:red;padding:20px;'>Lỗi đọc file: {exc}</div>", []

    if not isinstance(data, list):
        return "<div style='color:red;padding:20px;'>File JSON phải là một mảng (array) các object có trường \"text\".</div>", []

    model_keys = []
    if model_choice in ("GLiNER2", "Tất cả"):
        model_keys.append("multi_1")
    if model_choice in ("XLM-RoBERTa", "Tất cả"):
        model_keys.append("XLMBERT")
    if model_choice in ("mmBERT-multilingual", "Tất cả"):
        model_keys.append("mmBERT-multilingual")

    total = len(data)
    all_results = []
    html_parts = []

    # Summary header placeholder
    total_pii_gliner = 0
    total_pii_bert = 0
    total_pii_mmbert = 0

    for idx, item in enumerate(data):
        text = item.get("text", "")
        if not text:
            continue

        sample_num = idx + 1
        short_text = text[:150].replace("<", "&lt;").replace(">", "&gt;").replace("\n", " ")
        if len(text) > 150:
            short_text += "..."

        result_entry = {"sample": sample_num, "text_preview": text[:100]}

        # Run models
        html_gliner = preds_gliner = html_bert = preds_bert = html_mmbert = preds_mmbert = None
        if "multi_1" in model_keys:
            html_gliner, preds_gliner = run_single_model(text, "multi_1")
            result_entry["gliner"] = preds_gliner
            if isinstance(preds_gliner, list):
                total_pii_gliner += len(preds_gliner)
        if "XLMBERT" in model_keys:
            html_bert, preds_bert = run_single_model(text, "XLMBERT")
            result_entry["xlmbert"] = preds_bert
            if isinstance(preds_bert, list):
                total_pii_bert += len(preds_bert)
        if "mmBERT-multilingual" in model_keys:
            html_mmbert, preds_mmbert = run_single_model(text, "mmBERT-multilingual")
            result_entry["mmbert_multilingual"] = preds_mmbert
            if isinstance(preds_mmbert, list):
                total_pii_mmbert += len(preds_mmbert)

        all_results.append(result_entry)

        # Build sample HTML
        html_parts.append(f'<div class="batch-sample">')
        html_parts.append(
            f'<div class="batch-sample-header">'
            f'<span>📄 Mẫu #{sample_num} / {total}</span>'
            f'<span style="font-weight:normal;font-size:0.8rem;">{len(text)} ký tự</span>'
            f'</div>'
        )
        html_parts.append(f'<div class="batch-sample-input">{short_text}</div>')

        if len(model_keys) >= 2:
            cols = f"repeat({len(model_keys)}, 1fr)"
            html_parts.append(f'<div class="batch-models-row" style="grid-template-columns:{cols};">')
            _batch_cols = [
                ("multi_1", "gliner", "GLiNER2", html_gliner, preds_gliner),
                ("XLMBERT", "bert", "XLM-RoBERTa", html_bert, preds_bert),
                ("mmBERT-multilingual", "mmbert", "mmBERT-multi", html_mmbert, preds_mmbert),
            ]
            for mk, cls, name, h, p in _batch_cols:
                if mk not in model_keys:
                    continue
                html_parts.append('<div class="batch-model-col">')
                html_parts.append(f'<div class="batch-model-label {cls}">{name}</div>')
                if h:
                    html_parts.append(h)
                n = len(p) if isinstance(p, list) else 0
                html_parts.append(f'<div class="batch-entity-tags"><strong style="font-size:0.75rem;margin-right:6px;">{n} entities:</strong>')
                html_parts.append(_build_entity_tags_html(p))
                html_parts.append('</div></div>')
            html_parts.append('</div>')  # end batch-models-row
        else:
            # Single model
            mk = model_keys[0]
            _single_map = {
                "multi_1": ("gliner", "GLiNER2", html_gliner, preds_gliner),
                "XLMBERT": ("bert", "XLM-RoBERTa", html_bert, preds_bert),
                "mmBERT-multilingual": ("mmbert", "mmBERT-multi", html_mmbert, preds_mmbert),
            }
            cls, name, h, p = _single_map[mk]
            html_parts.append(f'<div style="padding:12px 16px;">')
            html_parts.append(f'<div class="batch-model-label {cls}">{name}</div>')
            if h:
                html_parts.append(h)
            n = len(p) if isinstance(p, list) else 0
            html_parts.append(f'<div class="batch-entity-tags"><strong style="font-size:0.75rem;margin-right:6px;">{n} entities:</strong>')
            html_parts.append(_build_entity_tags_html(p))
            html_parts.append('</div></div>')

        html_parts.append('</div>')  # end batch-sample

    # Build summary
    summary_lines = [
        f"<div class='batch-summary'>",
        f"<strong>📊 Tổng kết:</strong> {total} mẫu đã xử lý<br>",
    ]
    if "multi_1" in model_keys:
        summary_lines.append(f"🤖 <strong>GLiNER2:</strong> {total_pii_gliner} PII entities phát hiện<br>")
    if "XLMBERT" in model_keys:
        summary_lines.append(f"🧠 <strong>XLM-RoBERTa:</strong> {total_pii_bert} PII entities phát hiện<br>")
    if "mmBERT-multilingual" in model_keys:
        summary_lines.append(f"🔶 <strong>mmBERT-multilingual:</strong> {total_pii_mmbert} PII entities phát hiện")
    summary_lines.append("</div>")

    final_html = "".join(summary_lines) + "".join(html_parts)
    return final_html, all_results


# ==========================================
# 7. Giao diện Gradio UI — song song 2 model
# ==========================================
DEFAULT_TEXT = (
    "Dear Mr. PATEL,\n\n"
    "We are pleased to inform you that your application for a personal loan has been approved "
    "by RiverbankFinancial. The approval was finalized on 14-05-2024 09:45 and your documents "
    "will be processed within the next two business days. As a resident of Illinois, your "
    "application was reviewed in accordance with all regional regulations. Please note that your "
    "National ID, AID-6543217890, has been securely verified as part of our compliance process.\n\n"
    "For your records, your online application was submitted from the IP address "
    "2a02:4d60:1f31:4c3f:85e1:1122:abfc:0345. Your loan agreement and repayment schedule will "
    "be sent to you via email by john.patel@email.com. If you have any questions, please contact "
    "our support team at 555-0198.\n\n"
)

with gr.Blocks(title="PII Extractor — So sánh 3 Model", theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.Markdown("# 🔍 Hệ thống Trích xuất PII — So sánh 3 Model Song Song")

    with gr.Tabs():
        # ── Tab 1: Nhập văn bản ──────────────────────────────────────
        with gr.TabItem("📝 Nhập văn bản"):
            gr.Markdown(
                "Nhập văn bản bên dưới và nhấn **Trích xuất PII**. "
                "Kết quả của **GLiNER2 multi-v7**, **XLM-RoBERTa** và **mmBERT-multilingual** sẽ hiển thị đồng thời."
            )

            with gr.Row():
                input_text = gr.Textbox(
                    lines=12,
                    label="📝 Văn bản đầu vào",
                    value=DEFAULT_TEXT,
                    scale=1,
                )

            submit_btn = gr.Button("⚡ Trích xuất PII (Cả 3 Model)", variant="primary", size="lg")

            gr.Markdown("## 📊 Kết quả nhận diện")
            with gr.Row(equal_height=True):
                with gr.Column():
                    gr.HTML(
                        '<div class="model-header-gliner">🤖 Model 1 — GLiNER2 multi-v7 (checkpoint-epoch-3, threshold=0.7)</div>'
                    )
                    out_html_gliner = gr.HTML(label="GLiNER2 Output")

                with gr.Column():
                    gr.HTML(
                        '<div class="model-header-xlmbert">🧠 Model 2 — XLM-RoBERTa (FacebookAI, threshold=0.5)</div>'
                    )
                    out_html_bert = gr.HTML(label="XLM-RoBERTa Output")

                with gr.Column():
                    gr.HTML(
                        '<div class="model-header-mmbert">🔶 Model 3 — mmBERT-multilingual (jhu-clsp, threshold=0.5)</div>'
                    )
                    out_html_mmbert = gr.HTML(label="mmBERT-multilingual Output")

            gr.Markdown("## 🗂️ Dữ liệu JSON thô")
            with gr.Row(equal_height=True):
                with gr.Column():
                    gr.Markdown("**GLiNER2 JSON**")
                    out_json_gliner = gr.JSON(label="GLiNER2 Raw JSON")

                with gr.Column():
                    gr.Markdown("**XLM-RoBERTa JSON**")
                    out_json_bert = gr.JSON(label="XLM-RoBERTa Raw JSON")

                with gr.Column():
                    gr.Markdown("**mmBERT-multilingual JSON**")
                    out_json_mmbert = gr.JSON(label="mmBERT-multilingual Raw JSON")

            submit_btn.click(
                fn=process_both,
                inputs=[input_text],
                outputs=[out_html_gliner, out_json_gliner, out_html_bert, out_json_bert, out_html_mmbert, out_json_mmbert],
            )

        # ── Tab 2: Test File ─────────────────────────────────────────
        with gr.TabItem("📂 Test File"):
            gr.Markdown(
                "Upload file JSON chứa các mẫu test (format: `[{\"text\": \"...\"}, ...]`). "
                "Hệ thống sẽ chạy model trên từng mẫu và hiển thị kết quả trực quan."
            )

            with gr.Row():
                with gr.Column(scale=2):
                    test_file_input = gr.File(
                        label="📁 Upload file JSON test",
                        file_types=[".json"],
                        type="filepath",
                    )
                with gr.Column(scale=1):
                    model_choice = gr.Radio(
                        choices=["Tất cả", "GLiNER2", "XLM-RoBERTa", "mmBERT-multilingual"],
                        value="Tất cả",
                        label="🔧 Chọn model",
                    )

            test_btn = gr.Button("🚀 Chạy Test Batch", variant="primary", size="lg")

            gr.Markdown("## 📊 Kết quả Test")
            test_output_html = gr.HTML(label="Kết quả trực quan")
            with gr.Accordion("🗂️ JSON chi tiết", open=False):
                test_output_json = gr.JSON(label="Kết quả JSON")

            test_btn.click(
                fn=process_test_file,
                inputs=[test_file_input, model_choice],
                outputs=[test_output_html, test_output_json],
            )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)