#!/usr/bin/env python3
"""
PII Masking ZH-TW — inference module

Detects and masks 8 types of Taiwan-specific personal information (PII)
using a fine-tuned Qwen2.5-1.5B-Instruct + LoRA adapter, returning structured
JSON output.

Supported types: NAME, ROC_ID, PHONE, EMAIL, ADDRESS, NHI_ID, BANK_ACCOUNT,
CREDIT_CARD.

Quick start:
    from inference import PIIDetector

    detector = PIIDetector()
    result = detector.detect_and_replace(
        "客戶王小明來電諮詢，身分證A123456789，手機0912345678"
    )
    print(result["masked_text"])

Environment variables (optional, for local development):
    PII_BASE_MODEL    Override base model id or path
    PII_ADAPTER_REPO  Override adapter repo id or local path
"""

import json
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# ============================================================
# Configuration
# ============================================================
BASE_MODEL = os.environ.get("PII_BASE_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
ADAPTER_REPO = os.environ.get("PII_ADAPTER_REPO", "GOSHUNCLE/pii-masking-zh-tw")

SYSTEM_PROMPT = """你是一個台灣個人資料（PII）偵測專家。分析輸入文本，找出所有個人可識別資訊（PII），以嚴格 JSON 格式回傳。

偵測類型與遮罩規則：
- ROC_ID：身分證字號（格式：1英文字母+9數字）→ 保留首2碼與末2碼，其餘以*替代
- NAME：姓名（依語意判斷，或依特徵：2-4個中文字且首字為常見姓氏）→ 保留姓氏，名字以O替代
- PHONE：手機號碼（09開頭，10碼）→ 保留前4碼，第5碼起全部以*替代
- EMAIL：電子郵件 → 保留首2字元與@後域名，中間以*替代
- ADDRESS：完整地址（須包含「號」或「樓」才視為完整地址）→ 保留至「區」，區之後以*替代
- NHI_ID：健保卡號（12碼數字）→ 保留前4碼，其餘以*替代
- BANK_ACCOUNT：銀行帳號 → 保留前4碼，其餘以*替代
- CREDIT_CARD：信用卡號（16碼）→ 保留前4碼，其餘以*替代

輸出格式：
{"pii_found": true/false, "entities": [{"type": "類型", "value": "原始值", "masked": "遮罩值"}]}

規則：
1. 無PII時回傳 {"pii_found": false, "entities": []}
2. 同一文本中有多個PII須全部列出
3. 只回傳JSON，不加任何說明文字"""


# ============================================================
# JSON parsing (minimal defensive)
# ============================================================
def _extract_json(text):
    text = text.strip()
    start = text.find("{")
    end = text.rfind("}") + 1
    if start < 0 or end <= start:
        return None
    try:
        return json.loads(text[start:end])
    except json.JSONDecodeError:
        return None


# ============================================================
# Digit-character detection (halfwidth / fullwidth / Chinese numerals)
# ============================================================
_CHINESE_DIGITS = "零一二三四五六七八九"


def _is_digit_char(c):
    return c.isdigit() or c in _CHINESE_DIGITS


# ============================================================
# Filter 1: structural validation
# ------------------------------------------------------------
# Verify the entity's value matches its claimed type's structural pattern.
# Filters out misclassifications produced by the model.
# ============================================================
def _is_valid_entity(entity):
    etype = entity.get("type", "")
    value = entity.get("value", "")
    if not value:
        return False

    if etype == "NAME":
        return 2 <= len(value) <= 4

    if etype == "EMAIL":
        if "@" not in value:
            return False
        _, _, domain = value.partition("@")
        return "." in domain

    if etype == "ADDRESS":
        # System prompt requires "號" or "樓" for valid full address
        return "號" in value or "樓" in value

    raw = value.replace("-", "").replace(" ", "")

    if etype == "ROC_ID":
        if len(raw) != 10:
            return False
        return raw[0].isalpha() and all(_is_digit_char(c) for c in raw[1:])

    if etype == "PHONE":
        return len(raw) == 10 and all(_is_digit_char(c) for c in raw)

    if etype == "NHI_ID":
        return len(raw) == 12 and all(_is_digit_char(c) for c in raw)

    if etype == "BANK_ACCOUNT":
        return 10 <= len(raw) <= 16 and all(_is_digit_char(c) for c in raw)

    if etype == "CREDIT_CARD":
        return len(raw) == 16 and all(_is_digit_char(c) for c in raw)

    return True


# ============================================================
# Filter 2: provenance check (guards against hallucination)
# ------------------------------------------------------------
# The value must verbatim appear in the input text. Tolerant of
# whitespace/separator differences.
# ============================================================
def _value_in_text(value, text):
    if not value:
        return False
    if value in text:
        return True
    norm = lambda s: s.replace(" ", "").replace("-", "").replace("　", "")
    return norm(value) in norm(text)


# ============================================================
# Mask validation and correction
# ============================================================
def _validate_and_fix_mask(entity):
    """Recompute the masked field from the value field using deterministic
    rules. Compensates for the model's occasional miscount in star quantity."""
    etype = entity.get("type", "")
    value = entity.get("value", "")
    masked = entity.get("masked", "")

    correct_mask = None

    if etype == "ROC_ID" and len(value) == 10:
        correct_mask = value[:2] + "*" * 6 + value[-2:]

    elif etype == "NAME" and len(value) >= 2:
        correct_mask = value[0] + "O" * (len(value) - 1)

    elif etype == "PHONE":
        raw = value.replace("-", "").replace(" ", "")
        if len(raw) == 10:
            correct_mask = raw[:4] + "*" * 6

    elif etype == "EMAIL" and "@" in value:
        local, domain = value.split("@", 1)
        if len(local) >= 2:
            correct_mask = local[:2] + "*" * (len(local) - 2) + "@" + domain

    elif etype == "ADDRESS":
        # Find the deepest admin division as the cut point.
        # Priority: 區 / 鎮 / 鄉 first, fallback to 市 / 縣.
        prefix = None
        best_idx = -1
        for keyword in ["區", "鎮", "鄉"]:
            idx = value.find(keyword)
            if idx >= 0 and idx > best_idx:
                best_idx = idx
                prefix = value[:idx + 1]
        if prefix is None:
            for keyword in ["市", "縣"]:
                idx = value.find(keyword)
                if idx >= 0 and idx > best_idx:
                    best_idx = idx
                    prefix = value[:idx + 1]
        # Guard: prefix must be strictly shorter than value
        # (avoids cases like "社區" where 區 is at the end)
        if prefix and len(prefix) < len(value):
            correct_mask = prefix + "*" * (len(value) - len(prefix))

    elif etype in ("NHI_ID", "BANK_ACCOUNT", "CREDIT_CARD"):
        if len(value) >= 4:
            correct_mask = value[:4] + "*" * (len(value) - 4)

    if correct_mask and correct_mask != masked:
        entity["masked"] = correct_mask
        entity["mask_corrected"] = True
    return entity


# ============================================================
# PIIDetector
# ============================================================
class PIIDetector:
    def __init__(self, base_model=BASE_MODEL, adapter_repo=ADAPTER_REPO):
        print("Loading model...")
        self.tokenizer = AutoTokenizer.from_pretrained(
            base_model,
            trust_remote_code=True,
            padding_side="left",
        )
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        base = AutoModelForCausalLM.from_pretrained(
            base_model,
            dtype=torch.float32,
            device_map="cpu",
            trust_remote_code=True,
        )
        self.model = PeftModel.from_pretrained(base, adapter_repo)
        self.model.eval()
        print("Model loaded.")

    def detect(self, text, fix_masks=True):
        """Detect PII in `text`.

        Pipeline:
            1. Model inference
            2. JSON parse
            3. Filter 1: structural validation
            4. Filter 2: provenance check (rejects hallucinations)
            5. validate_and_fix_mask: deterministic mask recomputation
            6. Recompute pii_found (False if all entities filtered out)

        Returns:
            dict: {"pii_found": bool, "entities": list, "raw_response": str,
                   "parse_error": bool}
        """
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": text},
        ]
        prompt = self.tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
        )
        inputs = self.tokenizer(prompt, return_tensors="pt")

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=False,
                temperature=1.0,
                pad_token_id=self.tokenizer.pad_token_id,
            )
        new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        raw_response = self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

        parsed = _extract_json(raw_response)
        if parsed is None:
            return {
                "pii_found": False,
                "entities": [],
                "raw_response": raw_response,
                "parse_error": True,
            }

        # ===== Two-stage filtering =====
        entities = parsed.get("entities", [])
        if entities:
            # Filter 1: structural validation
            entities = [e for e in entities if _is_valid_entity(e)]
            # Filter 2: provenance check (rejects hallucinated/translated values)
            entities = [e for e in entities if _value_in_text(e.get("value", ""), text)]
            # Mask correction (only for entities passing both filters)
            if fix_masks:
                entities = [_validate_and_fix_mask(e) for e in entities]

        parsed["entities"] = entities
        parsed["pii_found"] = bool(entities)
        parsed["raw_response"] = raw_response
        parsed["parse_error"] = False
        return parsed

    def detect_and_replace(self, text):
        """Detect PII and return the input text with all detected PII replaced
        by their masked forms."""
        result = self.detect(text)
        masked_text = text

        if result.get("entities"):
            sorted_ents = sorted(
                result["entities"],
                key=lambda e: len(e.get("value", "")),
                reverse=True,
            )
            for e in sorted_ents:
                value = e.get("value", "")
                masked = e.get("masked", "")
                if value and masked and value in masked_text:
                    # Replace all occurrences
                    masked_text = masked_text.replace(value, masked)

        return {
            "original": text,
            "masked_text": masked_text,
            "entities": result.get("entities", []),
            "pii_found": result.get("pii_found", False),
        }


# ============================================================
# CLI
# ============================================================
if __name__ == "__main__":
    detector = PIIDetector()

    print("\n" + "=" * 60)
    print("PII Masking ZH-TW — interactive mode")
    print("Type text and press Enter to detect; type 'q' to exit.")
    print("=" * 60)

    while True:
        print()
        try:
            user_input = input("> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye.")
            break

        if user_input.lower() in ("q", "quit", "exit"):
            print("Bye.")
            break
        if not user_input:
            continue

        result = detector.detect_and_replace(user_input)

        print(f"\n  original: {result['original']}")
        print(f"  masked  : {result['masked_text']}")

        if result["entities"]:
            print(f"  found {len(result['entities'])} PII:")
            for e in result["entities"]:
                corrected = " (corrected)" if e.get("mask_corrected") else ""
                print(f"    [{e['type']}] {e['value']} -> {e['masked']}{corrected}")
        else:
            print("  no PII detected")