import asyncio
import atexit
import json
import os
import re
import tempfile
import unicodedata
from functools import lru_cache
from pathlib import Path
from typing import Dict, List, Optional, Tuple


def _close_event_loop() -> None:
    try:
        loop = asyncio.get_event_loop()
        if not loop.is_closed():
            loop.close()
    except Exception:
        pass


atexit.register(_close_event_loop)

import gradio as gr

try:
    from huggingface_hub import login as hf_login
except Exception:
    hf_login = None

try:
    import torch
except Exception:
    torch = None

from gliner2 import GLiNER2
from rule_detector import DEFAULT_RULEBASE_ALLOWED_ENTITY_TYPES, detect_by_rules


DEFAULT_MODEL_ID = os.getenv("MODEL_ID", "AITeamUIT/gliner2-multi-v1-e3-25-6")
DEFAULT_THRESHOLD = float(os.getenv("THRESHOLD", "0.5"))
DEFAULT_CHUNK_CHARS = int(os.getenv("CHUNK_CHARS", "1000"))
DEFAULT_CHUNK_TOKENS = int(os.getenv("CHUNK_TOKENS", "768"))
CHUNK_SAFETY_MARGIN = int(os.getenv("CHUNK_SAFETY_MARGIN", "20"))
MAX_CHARS_PER_CHUNK = int(os.getenv("MAX_CHARS_PER_CHUNK", "8000"))
MAX_WIDTH = int(os.getenv("MAX_WIDTH", "30"))

PII_LABEL_DESCRIPTIONS: Dict[str, str] = {
    "PERSON": "Person names",
    "DATE": "Calendar dates",
    "JOB_TITLE": "Job titles",
    "TIME": "Time values",
    "LOCATION": "General locations",
    "PREFIX": "Titles before names",
    "ORGANIZATION": "Organization names",
    "EMAIL": "Email addresses",
    "PHONE": "Phone numbers",
    "ADDRESS": "Street addresses",
    "GENDER": "Sex or gender",
    "COORDINATE": "Geographic coordinates",
    "MONEY": "Monetary amounts",
    "IP": "IP addresses",
    "ZIP_CODE": "Postal or ZIP codes",
    "NATIONALITY": "Legal nationality",
    "AGE": "Age values",
    "USERNAME": "Usernames or handles",
    "URL": "Web links or URLs",
    "MARITAL": "Marital status",
    "TRADE_UNION": "Trade union membership",
    "BIRTHDATE": "Date of birth",
    "CARD_ISSUER": "Card brands or issuers",
    "PIN": "Authentication PINs",
    "PASSPORT": "Passport numbers",
    "MEDICAL_INFO": "Health or medical information",
    "IBAN": "International bank account numbers",
    "PASSWORD": "Account passwords",
    "ETHNICITY": "Ethnic or racial group",
    "RELIGION": "Religious affiliation",
    "EMPLOYEE_ID": "Employee identifiers",
    "INSURANCE_ID": "Insurance identifiers",
    "TIN": "Tax identification numbers",
    "SWIFT": "Bank SWIFT or BIC codes",
    "CARD_NUMBER": "Payment card numbers",
    "ACCOUNT_ID": "Account identifiers",
    "WALLET": "Crypto wallet addresses",
    "NATIONAL_ID": "National ID numbers",
    "CVV": "Card security codes",
    "BANK_ACCOUNT": "Bank account numbers",
    "LICENSE": "Driver license numbers",
    "TICKET_ID": "Support or case IDs",
    "PLATE": "Vehicle plate numbers",
    "API_KEY": "API keys or tokens",
}

LABELS = list(PII_LABEL_DESCRIPTIONS.keys())
RULE_LABELS = set(DEFAULT_RULEBASE_ALLOWED_ENTITY_TYPES)

LABEL_PRIORITY: Dict[str, int] = {
    "WALLET": 10,
    "IBAN": 10,
    "CARD_NUMBER": 10,
    "NATIONAL_ID": 10,
    "IP": 10,
    "URL": 10,
    "EMAIL": 10,
    "MONEY": 10,
    "AGE": 10,
}

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]+")
NUMERIC_MASK_LABELS = {
    "CARD_NUMBER",
    "PHONE",
    "PIN",
    "CVV",
    "TIN",
    "BANK_ACCOUNT",
    "IBAN",
    "SWIFT",
    "WALLET",
    "IP",
    "ZIP_CODE",
}


try:
    from gliner2.processor import WhitespaceTokenSplitter

    IMPROVED_PATTERN = re.compile(
        r"""
        # Scheme / www URL
        (?:(?:https?|ftp)://|www\.)[^\s<>"'`{}|\\^]*[^\s<>"'`{}|\\^.,;:!?()\[\]]
    
        # Email address
        |[a-z0-9.!#$%&'*+/=?^_`{|}~-]+@(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)+[a-z]{2,63}
    
        # Social/user handle
        |@[a-z0-9_](?:[a-z0-9_.-]{0,28}[a-z0-9_])?
    
        # IPv4 address. Allows sentence punctuation after the IP, but avoids
        # matching inside a longer dotted numeric sequence.
        |(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)(?:\.(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)){3}(?=$|[^\w])(?!\.\d)
    
        # IPv6 address, including compressed forms and optional zone IDs.
        |(?:
            (?:[0-9a-f]{1,4}:){7}[0-9a-f]{1,4}
            |[0-9a-f]{1,4}:(?:(?::[0-9a-f]{1,4}){1,6})
            |(?:[0-9a-f]{1,4}:){1,2}(?::[0-9a-f]{1,4}){1,5}
            |(?:[0-9a-f]{1,4}:){1,3}(?::[0-9a-f]{1,4}){1,4}
            |(?:[0-9a-f]{1,4}:){1,4}(?::[0-9a-f]{1,4}){1,3}
            |(?:[0-9a-f]{1,4}:){1,5}(?::[0-9a-f]{1,4}){1,2}
            |(?:[0-9a-f]{1,4}:){1,6}:[0-9a-f]{1,4}
            |:(?:(?::[0-9a-f]{1,4}){1,7}|:)
            |(?:[0-9a-f]{1,4}:){1,7}:
        )(?:%[0-9a-z_.-]+)?(?![0-9a-f:])
    
        # Strict JWT-like token
        |[a-z0-9_-]{10,}\.[a-z0-9_-]{10,}\.[a-z0-9_-]{10,}
    
        # Shorter dot-separated token-like string, e.g. sk-abc.def.ghi
        |(?=[a-z0-9_.-]*[-_\d])[a-z0-9_-]{2,}(?:\.[a-z0-9_-]{2,}){2,}
    
        # Bare domain / URL-like candidate. Validation belongs downstream; this
        # tokenizer only keeps plausible URL spans from being split into pieces.
        |(?<![\w@])
            (?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)+
            (?:[a-z]{2,63}|xn--[a-z0-9-]{2,59})
            (?=$|/|:|[^\w-])
            (?::\d{2,5})?
            (?:/(?:[^\s<>"'`{}|\\^]*[^\s<>"'`{}|\\^.,;:!?()\[\]])?)?
    
        # Letter-containing word token: must have at least one letter or
        # underscore (the [^\W\d] anchor). Only underscores join tokens;
        # hyphens always split (except digit-digit handled below).
        |(?:\d*[^\W\d]\w*)(?:_\w+)*
    
        # Pure digit sequence: hyphens/underscores only join other digit groups.
        # "555-1234" stays atomic; "13-tuổi" → "13" / "-" / "tuổi" so the
        # numeric entity (AGE, DATE, …) is not fused with the surrounding word.
        |\d+(?:[-_]\d+)*
    
        # Fallback: any single non-space character.
        |\S
        """,
        re.VERBOSE | re.IGNORECASE,
    )
    WhitespaceTokenSplitter._PATTERN = IMPROVED_PATTERN
except Exception:
    pass


@lru_cache(maxsize=2)
def load_model(model_id: str):
    hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
    if hf_token and hf_login is not None:
        hf_login(token=hf_token)
    model = GLiNER2.from_pretrained(model_id)
    model.config.max_width = MAX_WIDTH
    model.max_width = MAX_WIDTH
    if hasattr(model, "span_rep") and hasattr(model.span_rep, "span_rep_layer"):
        model.span_rep.span_rep_layer.max_width = MAX_WIDTH
    if torch is not None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model = model.to(device)
    model.eval()
    return model


def build_schema(model, use_descriptions: bool):
    schema = model.create_schema()
    if use_descriptions:
        schema.entities(PII_LABEL_DESCRIPTIONS)
    else:
        schema.entities(LABELS)
    return schema


def manual_chunks(text: str, max_chars: int) -> List[Tuple[str, int]]:
    if len(text) <= max_chars:
        return [(text, 0)]
    tokens = list(re.finditer(r"\S+", text))
    if not tokens:
        return [(text, 0)]
    chunks: List[Tuple[str, int]] = []
    chunk_start = 0
    for index in range(1, len(tokens)):
        seg_start = tokens[chunk_start].start()
        seg_end = tokens[index].end()
        if seg_end - seg_start > max_chars:
            prev_end = tokens[index - 1].end()
            chunks.append((text[seg_start:prev_end], seg_start))
            chunk_start = index
    seg_start = tokens[chunk_start].start()
    chunks.append((text[seg_start:tokens[-1].end()], seg_start))
    return chunks


def word_token_count(text: str) -> int:
    """Count word tokens the way the (patched) model tokenizer splits them."""
    try:
        from gliner2.processor import WhitespaceTokenSplitter

        splitter = WhitespaceTokenSplitter()
        return sum(1 for _ in splitter(text, lower=False))
    except Exception:
        return len(text.split())


class RecursiveWordChunker:
    """Label-aware recursive chunker built on chonkie.RecursiveChunker.

    The chunk-size budget is measured in the model's own word-token space and
    already subtracts the entity-schema label tokens, so each emitted chunk plus
    the labels fits the encoder. Recursion falls back paragraph -> sentence ->
    whitespace. ``chunk()`` returns ``(chunk_text, start_index)`` tuples.
    """

    def __init__(self, budget_words: int, max_chars: int = MAX_CHARS_PER_CHUNK, min_chars_per_chunk: int = 24):
        # chonkie is a pip dependency (not a cross-folder import).
        from chonkie.chunker.recursive import RecursiveChunker
        from chonkie.types import RecursiveLevel, RecursiveRules

        self.budget = budget_words
        self.max_chars = max_chars
        self._chunker = RecursiveChunker(
            tokenizer=word_token_count,
            chunk_size=budget_words,
            rules=RecursiveRules(
                levels=[
                    RecursiveLevel(delimiters=["\n\n", "\r\n", "\n", "\r"]),
                    RecursiveLevel(delimiters=[". ", "! ", "? ", ".\n", "!\n", "?\n"]),
                    RecursiveLevel(whitespace=True),
                ]
            ),
            min_characters_per_chunk=min_chars_per_chunk,
        )

    def chunk(self, text: str) -> List[Tuple[str, int]]:
        return [(c.text, int(c.start_index)) for c in self._chunker.chunk(text)]


@lru_cache(maxsize=4)
def build_word_chunker(
    chunk_tokens: int = DEFAULT_CHUNK_TOKENS,
    safety_margin: int = CHUNK_SAFETY_MARGIN,
    max_chars: int = MAX_CHARS_PER_CHUNK,
) -> RecursiveWordChunker:
    """Build a :class:`RecursiveWordChunker` with a label-aware token budget.

    The label tokens (``"<LABEL> <description>"`` plus a +2 separator budget each)
    are subtracted from ``chunk_tokens`` so the encoder never overflows.
    """
    labels_words = [f"{k} {v}" for k, v in PII_LABEL_DESCRIPTIONS.items()]
    label_token_count = sum(word_token_count(lbl) + 2 for lbl in labels_words)
    budget = max(64, chunk_tokens - label_token_count - safety_margin)
    return RecursiveWordChunker(budget_words=budget, max_chars=max_chars)


def chunk_text(text: str) -> List[Tuple[str, int]]:
    """Split text into (chunk_text, start_index) using the recursive word chunker.

    Falls back to a plain whitespace chunker only if chonkie is unavailable.
    """
    try:
        return build_word_chunker(DEFAULT_CHUNK_TOKENS).chunk(text)
    except Exception:
        return manual_chunks(text, DEFAULT_CHUNK_CHARS)
    spans = []
    if not value:
        return spans
    start = 0
    while True:
        index = text.find(value, start)
        if index == -1:
            break
        spans.append((index, index + len(value)))
        start = index + max(1, len(value))
    return spans


def normalize_predictions(raw_preds, text: str, emit_duplicates: bool = False) -> List[Dict]:
    spans = []
    if isinstance(raw_preds, dict):
        entities = raw_preds.get("entities", raw_preds)
        if isinstance(entities, list):
            entities = entities[0] if entities and isinstance(entities[0], dict) else {}
        if not isinstance(entities, dict):
            entities = {}
        for label, items in entities.items():
            if not isinstance(items, list):
                continue
            label = str(label).upper().strip()
            for item in items:
                if isinstance(item, dict):
                    start = item.get("start")
                    end = item.get("end")
                    value = item.get("text") or item.get("value") or ""
                    score = float(item.get("score", item.get("confidence", 1.0)) or 0.0)
                    has_span = start is not None and end is not None and int(end) > int(start)
                    if emit_duplicates and value:
                        for match_start, match_end in find_all_occurrences(text, str(value)):
                            spans.append({"label": label, "start": match_start, "end": match_end, "text": str(value), "score": score})
                    elif has_span:
                        start_int = int(start)
                        end_int = int(end)
                        spans.append({"label": label, "start": start_int, "end": end_int, "text": item.get("text", text[start_int:end_int]), "score": score})
                    elif value:
                        for match_start, match_end in find_all_occurrences(text, str(value)):
                            spans.append({"label": label, "start": match_start, "end": match_end, "text": str(value), "score": score})
                elif isinstance(item, str):
                    for match_start, match_end in find_all_occurrences(text, item):
                        spans.append({"label": label, "start": match_start, "end": match_end, "text": item, "score": 0.0})
    elif isinstance(raw_preds, list):
        spans = list(raw_preds)

    best: Dict[tuple, Dict] = {}
    for span in spans:
        try:
            key = (span.get("label"), int(span.get("start")), int(span.get("end")))
        except Exception:
            continue
        if key not in best or float(span.get("score", 0) or 0) > float(best[key].get("score", 0) or 0):
            span["label"] = str(span.get("label", "")).upper().strip()
            span["start"] = key[1]
            span["end"] = key[2]
            span["text"] = str(span.get("text") or text[key[1]:key[2]])
            span["score"] = float(span.get("score", 0) or 0)
            best[key] = span
    return list(best.values())


def clean_text_with_mapping(text: str) -> Tuple[str, List[int]]:
    if not text:
        return text, []
    original = unicodedata.normalize("NFC", text)
    cleaned_chars: List[str] = []
    char_map: List[int] = []
    for original_pos, char in enumerate(original):
        if char in STRIP_CHARS or C0_RE.match(char):
            continue
        if C1_RE.match(char):
            cleaned_chars.append(" ")
            char_map.append(original_pos)
            continue
        cleaned_chars.append(char)
        char_map.append(original_pos)

    cleaned = MULTI_SPACE_RE.sub(" ", "".join(cleaned_chars)).strip()
    final_map: List[int] = []
    map_index = 0
    for char in cleaned:
        while map_index < len(char_map) and (
            original[char_map[map_index]] != char
            and not (char == " " and original[char_map[map_index]].isspace())
        ):
            map_index += 1
        final_map.append(char_map[map_index] if map_index < len(char_map) else len(text))
        map_index += 1
    final_map.append(len(text))
    return cleaned, final_map


def spans_overlap(left: Dict, right: Dict) -> bool:
    return max(int(left["start"]), int(right["start"])) < min(int(left["end"]), int(right["end"]))


def resolve_span_conflicts(entities: List[Dict], strategy: str = "score_first") -> List[Dict]:
    if not entities:
        return []

    def sort_key(entity: Dict):
        is_rule = 1 if entity.get("source") == "rule" else 0
        score = 1.0 if is_rule else float(entity.get("score") or 0.0)
        length = int(entity["end"]) - int(entity["start"])
        priority = LABEL_PRIORITY.get(str(entity.get("label", "")), 0)
        if strategy == "rule_first":
            return (is_rule, priority, score, length)
        if strategy == "longest_first":
            return (is_rule, length, score, priority)
        return (is_rule, score, priority, length)

    ranked = sorted(entities, key=sort_key, reverse=True)
    kept: List[Dict] = []
    for entity in ranked:
        if not any(spans_overlap(entity, existing) for existing in kept):
            kept.append(entity)
    return sorted(kept, key=lambda item: int(item["start"]))


def apply_span_validators(entities: List[Dict]) -> List[Dict]:
    kept = []
    for entity in entities:
        value = str(entity.get("text") or "")
        label = str(entity.get("label") or "").upper()
        if re.search(r"\*{2,}", value):
            continue
        if label in NUMERIC_MASK_LABELS and re.search(r"X{2,}", value):
            continue
        kept.append(entity)
    return kept


def apply_postprocessing(entities: List[Dict], original_text: str) -> List[Dict]:
    entities = apply_span_validators(entities)
    sorted_entities = sorted(entities, key=lambda item: int(item["start"]))
    final_entities: List[Dict] = []
    for index, entity in enumerate(sorted_entities):
        label = entity.get("label")
        start = int(entity["start"])
        end = int(entity["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 index + 1 < len(sorted_entities):
                next_entity = sorted_entities[index + 1]
                gap = int(next_entity["start"]) - end
                if next_entity.get("label") == "PERSON" and 0 <= gap <= 3:
                    has_person_after = True
            if not has_person_after:
                continue
        final_entities.append(entity)
    return final_entities


def tokenize_chunk(text: str) -> List[Tuple[str, int, int]]:
    """Run WhitespaceTokenSplitter on text, return list of (token, start, end)."""
    try:
        from gliner2.processor import WhitespaceTokenSplitter
        splitter = WhitespaceTokenSplitter()
        return list(splitter(text, lower=False))
    except Exception:
        return []


def run_inference(text: str, model_id: str, threshold: float, use_rules: bool, use_descriptions: bool, emit_duplicates: bool) -> Tuple[List[Dict], List[Dict], Dict]:
    """Returns (resolved_entities, raw_model_entities, debug_info)."""
    text = text.rstrip()
    debug_info = {"chunks": [], "tokens_per_chunk": [], "raw_preds_per_chunk": [], "token_merge_info": []}
    if not text or not text.strip():
        return [], [], debug_info
    model = load_model(model_id)
    schema = build_schema(model, use_descriptions)
    all_predictions: List[Dict] = []

    chunks = chunk_text(text)
    for chunk_idx, (chunk_text_value, chunk_start) in enumerate(chunks):
        # Store chunk info
        debug_info["chunks"].append({"index": chunk_idx, "start": chunk_start, "end": chunk_start + len(chunk_text_value), "text": chunk_text_value})

        # Tokenize chunk (show what WhitespaceTokenSplitter produces)
        tokens = tokenize_chunk(chunk_text_value)
        debug_info["tokens_per_chunk"].append({"chunk_index": chunk_idx, "tokens": [(t, s, e) for t, s, e in tokens]})

        raw = model.extract(
            chunk_text_value,
            schema,
            threshold=threshold,
            include_spans=not emit_duplicates,
            include_confidence=not emit_duplicates,
            format_results=emit_duplicates,
        )
        predictions = normalize_predictions(raw, chunk_text_value, emit_duplicates=emit_duplicates)

        # Store per-chunk raw predictions (before offset adjustment)
        debug_info["raw_preds_per_chunk"].append({"chunk_index": chunk_idx, "predictions": [dict(p) for p in predictions]})

        # Build token merge info: show how predicted spans map to tokens
        for pred in predictions:
            covered_tokens = [(t, s, e) for t, s, e in tokens if s < pred["end"] and e > pred["start"]]
            debug_info["token_merge_info"].append({
                "chunk_index": chunk_idx,
                "label": pred.get("label", ""),
                "text": pred.get("text", ""),
                "start": pred["start"],
                "end": pred["end"],
                "score": pred.get("score", 0.0),
                "tokens_merged": covered_tokens,
            })

        for prediction in predictions:
            prediction["start"] += chunk_start
            prediction["end"] += chunk_start
        all_predictions.extend(predictions)

    # Capture all raw model predictions before overlap resolution
    raw_model_entities = [dict(p, source="model") for p in all_predictions]

    if use_rules:
        cleaned, char_map = clean_text_with_mapping(text)
        original_to_clean: Dict[int, int] = {}
        for clean_index, original_pos in enumerate(char_map):
            original_to_clean.setdefault(original_pos, clean_index)

        model_predictions_clean = []
        for prediction in all_predictions:
            clean_start = original_to_clean.get(prediction["start"])
            clean_end = original_to_clean.get(prediction["end"])
            if clean_end is None:
                # p["end"] may land on a stripped char (e.g. trailing space).
                # Scan forward to the next mapped original position.
                for k in range(p["end"] + 1, len(text) + 2):
                    clean_end = orig_to_clean.get(k)
                    if clean_end is not None:
                        break
            if clean_start is None or clean_end is None or clean_end <= clean_start:
                continue
            model_predictions_clean.append({
                "label": prediction["label"],
                "text": cleaned[clean_start:clean_end],
                "start": clean_start,
                "end": clean_end,
                "score": prediction.get("score", 0.0),
                "source": "model",
            })

        rule_predictions = [
            {**entity, "score": 1.0, "source": "rule"}
            for entity in detect_by_rules(cleaned, LABELS)
        ]
        combined = resolve_span_conflicts(rule_predictions + model_predictions_clean, strategy="rule_first")
        restored = []
        for entity in combined:
            start = int(entity["start"])
            end = int(entity["end"])
            original_start = char_map[start] if start < len(char_map) else None
            original_end = char_map[end - 1] + 1 if (end - 1) < len(char_map) else None
            if original_start is None or original_end is None or original_end <= original_start:
                continue
            restored.append({
                "label": entity["label"],
                "start": original_start,
                "end": original_end,
                "text": text[original_start:original_end],
                "score": float(entity.get("score", 1.0)),
                "source": entity.get("source", "model"),
            })
        all_predictions = restored
    else:
        for prediction in all_predictions:
            prediction.setdefault("source", "model")
        all_predictions = resolve_span_conflicts(all_predictions)

    return apply_postprocessing(all_predictions, text), raw_model_entities, debug_info


def extract_with_llamaparse(path: str) -> List[str]:
    """Use LlamaParse (cloud) to parse a PDF or DOCX file — returns one string per document/page."""
    api_key = os.getenv("LLAMA_CLOUD_API_KEY") or os.getenv("LLAMA_API_KEY")
    if not api_key:
        raise RuntimeError("Set LLAMA_CLOUD_API_KEY in environment to use LlamaParse")

    try:
        from llama_parse import LlamaParse
    except Exception:
        try:
            from llama_cloud_services import LlamaParse
        except Exception as exc:
            raise RuntimeError("Install llama-parse to use LlamaParse parsing") from exc

    parser = LlamaParse(api_key=api_key, result_type="markdown", verbose=False)
    documents = parser.load_data(path)
    if not documents:
        return [""]
    return [getattr(document, "text", "") or "" for document in documents]


def extract_with_llamaindex_local(path: str, suffix: str) -> List[str]:
    """Use LlamaIndex local readers to parse a PDF or DOCX file — returns one string per page/document."""
    if suffix == ".docx":
        from llama_index.readers.file import DocxReader
        docs = DocxReader().load_data(Path(path))
        return ["\n".join(d.text for d in docs)] if docs else [""]

    if suffix == ".pdf":
        for cls_name in ("PyMuPDFReader", "PDFReader"):
            try:
                module = __import__("llama_index.readers.file", fromlist=[cls_name])
                reader_cls = getattr(module, cls_name)
                docs = reader_cls().load_data(Path(path))
                if any(d.text.strip() for d in docs):
                    return [d.text for d in docs]
            except Exception:
                continue
        return [""]

    raise ValueError(f"Unsupported file type: {suffix}")


def extract_input_text(text_input: str, file_obj, pdf_backend: str) -> Tuple[str, str, List[int]]:
    if file_obj is None:
        return text_input or "", "text", []

    file_path = Path(file_obj.name if hasattr(file_obj, "name") else str(file_obj))
    suffix = file_path.suffix.lower()
    if suffix not in (".docx", ".pdf"):
        raise ValueError("Only .docx and .pdf files are supported")

    if pdf_backend == "llamaparse":
        pages = extract_with_llamaparse(str(file_path))
    else:
        pages = extract_with_llamaindex_local(str(file_path), suffix)

    if suffix == ".docx":
        return "\n".join(pages), f"docx: {file_path.name}", []

    # PDF: merge pages with separators and track page offsets
    offsets: List[int] = []
    parts: List[str] = []
    cursor = 0
    for page_index, page_text in enumerate(pages, start=1):
        if page_index > 1:
            separator = "\n\n"
            parts.append(separator)
            cursor += len(separator)
        offsets.append(cursor)
        parts.append(page_text)
        cursor += len(page_text)
    return "".join(parts), f"pdf: {file_path.name}", offsets


def add_page_numbers(entities: List[Dict], page_offsets: List[int]) -> List[Dict]:
    if not page_offsets:
        return entities
    enriched = []
    for entity in entities:
        page = 1
        for index, offset in enumerate(page_offsets, start=1):
            if entity["start"] >= offset:
                page = index
            else:
                break
        enriched.append({**entity, "page": page})
    return enriched


def entities_to_table(entities: List[Dict]) -> List[List]:
    rows = []
    for index, entity in enumerate(entities, start=1):
        rows.append([
            index,
            entity.get("label", ""),
            entity.get("text", ""),
            int(entity.get("start", 0)),
            int(entity.get("end", 0)),
            round(float(entity.get("score", 0.0)), 4),
            entity.get("source", "model"),
            entity.get("page", ""),
        ])
    return rows


def make_highlighted_text(text: str, entities: List[Dict]):
    if not text:
        return []
    valid_entities = sorted(
        [entity for entity in entities if 0 <= int(entity.get("start", -1)) < int(entity.get("end", -1)) <= len(text)],
        key=lambda item: int(item["start"]),
    )
    highlights = []
    cursor = 0
    for entity in valid_entities:
        start = int(entity["start"])
        end = int(entity["end"])
        if start > cursor:
            highlights.append((text[cursor:start], None))
        highlights.append((text[start:end], entity.get("label", "PII")))
        cursor = end
    if cursor < len(text):
        highlights.append((text[cursor:], None))
    return highlights


def write_json_result(payload: Dict) -> str:
    temp_dir = tempfile.mkdtemp(prefix="guardpii_demo_")
    output_path = os.path.join(temp_dir, "pii_predictions.json")
    with open(output_path, "w", encoding="utf-8") as output_file:
        json.dump(payload, output_file, ensure_ascii=False, indent=2)
    return output_path


def detect_pii(text_input, file_obj, model_id, threshold, use_rules, use_descriptions, emit_duplicates, pdf_backend):
    text, source, page_offsets = extract_input_text(text_input, file_obj, pdf_backend)
    if not text.strip():
        empty_payload = {"source": source, "num_entities": 0, "entities": [], "text": text}
        return [], [], empty_payload, write_json_result(empty_payload), "No text was extracted.", "", "", "", ""

    entities, raw_model_entities, debug_info = run_inference(
        text=text,
        model_id=(model_id or DEFAULT_MODEL_ID).strip(),
        threshold=float(threshold),
        use_rules=bool(use_rules),
        use_descriptions=bool(use_descriptions),
        emit_duplicates=bool(emit_duplicates),
    )
    entities = add_page_numbers(entities, page_offsets)
    raw_model_entities = add_page_numbers(raw_model_entities, page_offsets)
    payload = {
        "source": source,
        "num_chars": len(text),
        "num_entities": len(entities),
        "model": (model_id or DEFAULT_MODEL_ID).strip(),
        "threshold": float(threshold),
        "use_rules": bool(use_rules),
        "entities": entities,
        "text": text,
    }
    status = f"Found {len(entities)} entities in {source}."

    # --- Debug output 1: Chunks ---
    chunks_md = "## Chunks Created\n\n"
    for chunk in debug_info["chunks"]:
        chunks_md += f"### Chunk {chunk['index']} (chars {chunk['start']}–{chunk['end']}, len={chunk['end'] - chunk['start']})\n"
        preview = chunk["text"][:500] + ("..." if len(chunk["text"]) > 500 else "")
        chunks_md += f"```\n{preview}\n```\n\n"

    # --- Debug output 2: WhitespaceTokenSplitter output ---
    tokens_md = "## WhitespaceTokenSplitter Output\n\n"
    for item in debug_info["tokens_per_chunk"]:
        tokens_md += f"### Chunk {item['chunk_index']} — {len(item['tokens'])} tokens\n\n"
        tokens_md += "| # | Token | Start | End |\n|---|-------|-------|-----|\n"
        for idx, (tok, s, e) in enumerate(item["tokens"][:200]):
            tok_display = tok.replace("|", "\\|")
            tokens_md += f"| {idx} | `{tok_display}` | {s} | {e} |\n"
        if len(item["tokens"]) > 200:
            tokens_md += f"\n... và {len(item['tokens']) - 200} tokens nữa\n"
        tokens_md += "\n"

    # --- Debug output 3: Token merge info (how model predictions map to tokens) ---
    merge_md = "## Token → Prediction Merge\n\n"
    merge_md += "Hiển thị cách model predict trên token spans và ghép lại thành entity:\n\n"
    for info in debug_info["token_merge_info"]:
        tokens_str = " + ".join([f"`{t}`" for t, s, e in info["tokens_merged"]])
        merge_md += f"- **[{info['label']}]** \"{info['text']}\" (chunk {info['chunk_index']}, chars {info['start']}–{info['end']}, score={info['score']:.4f})\n"
        merge_md += f"  - Tokens merged: {tokens_str}\n\n"

    # --- Debug output 4: Per-chunk predictions ---
    per_chunk_md = "## Per-Chunk Model Predictions\n\n"
    for item in debug_info["raw_preds_per_chunk"]:
        per_chunk_md += f"### Chunk {item['chunk_index']} — {len(item['predictions'])} predictions\n\n"
        if item["predictions"]:
            per_chunk_md += "| # | Label | Text | Start | End | Score |\n|---|-------|------|-------|-----|-------|\n"
            for idx, pred in enumerate(item["predictions"], 1):
                text_display = pred.get("text", "").replace("|", "\\|")[:80]
                per_chunk_md += f"| {idx} | {pred.get('label', '')} | {text_display} | {pred.get('start', '')} | {pred.get('end', '')} | {pred.get('score', 0.0):.4f} |\n"
            per_chunk_md += "\n"
        else:
            per_chunk_md += "Không có prediction nào.\n\n"

    return make_highlighted_text(text, entities), entities_to_table(raw_model_entities), payload, write_json_result(payload), status, chunks_md, tokens_md, merge_md, per_chunk_md


with gr.Blocks(title="GuardPII GLiNER2 Demo") as demo:
    gr.Markdown("# GuardPII GLiNER2 PII Demo")
    gr.Markdown("Detect PII from text, DOCX, or PDF files with GLiNER2 plus the high-precision rulebase.")

    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="Text input",
                lines=10,
                placeholder="Paste text here, or upload a DOCX/PDF below.",
            )
            file_input = gr.File(label="DOCX or PDF input", file_types=[".docx", ".pdf"])
            run_button = gr.Button("Detect PII", variant="primary")
        with gr.Column(scale=1):
            model_id = gr.Textbox(label="Model ID", value=DEFAULT_MODEL_ID)
            threshold = gr.Slider(label="Threshold", minimum=0.05, maximum=0.95, step=0.05, value=DEFAULT_THRESHOLD)
            use_rules = gr.Checkbox(label="Use rulebase gap-filler", value=True)
            use_descriptions = gr.Checkbox(label="Use label descriptions", value=False)
            emit_duplicates = gr.Checkbox(label="Value broadcast mode", value=False)
            pdf_backend = gr.Dropdown(
                label="Document parser",
                choices=["llamaparse", "llamaindex"],
                value=os.getenv("PDF_BACKEND", "llamaparse"),
            )

    status = gr.Textbox(label="Status", interactive=False)
    highlighted = gr.HighlightedText(label="Highlighted text", combine_adjacent=True, show_legend=True)
    table = gr.Dataframe(
        label="Entities",
        headers=["#", "label", "text", "start", "end", "score", "source", "page"],
        datatype=["number", "str", "str", "number", "number", "number", "str", "str"],
        wrap=True,
    )
    json_output = gr.JSON(label="JSON result")
    json_file = gr.File(label="Download JSON")

    gr.Markdown("---")
    gr.Markdown("# Debug / Pipeline Visualization")

    with gr.Accordion("1. Chunks (text splitting)", open=False):
        debug_chunks = gr.Markdown(value="")

    with gr.Accordion("2. WhitespaceTokenSplitter Output", open=False):
        debug_tokens = gr.Markdown(value="")

    with gr.Accordion("3. Token → Prediction Merge (how tokens are joined into entities)", open=False):
        debug_merge = gr.Markdown(value="")

    with gr.Accordion("4. Per-Chunk Model Predictions", open=False):
        debug_per_chunk = gr.Markdown(value="")

    run_button.click(
        detect_pii,
        inputs=[text_input, file_input, model_id, threshold, use_rules, use_descriptions, emit_duplicates, pdf_backend],
        outputs=[highlighted, table, json_output, json_file, status, debug_chunks, debug_tokens, debug_merge, debug_per_chunk],
    )


if __name__ == "__main__":
    demo.launch()