import json
from pathlib import Path

import numpy as np

LABELS = ["first_name", "last_name", "email", "phone_number", "street_address", "customer_id", "account_number"]
BIO = ["O"] + [f"{prefix}-{label}" for label in LABELS for prefix in ("B", "I")]
MODEL_REVISION = "6ea81172465e8b0ad3fddeed32b986cdcdcffcf0"


def byte_positions(text):
    positions = [0]
    for ch in text:
        positions.append(positions[-1] + len(ch.encode("utf-8")))
    return positions


def encode_windows(tokenizer, text, max_length, stride):
    """Tokenize the complete input before slicing; never use tokenizer overflow.

    tokenizers 0.23.2 early-exits before producing overflow, silently losing text.
    DistilBERT uses one CLS and one SEP token; raw offsets remain original chars.
    """
    capacity = max_length - 2
    if not 0 <= stride < capacity:
        raise ValueError("invalid stride or max_length")
    if tokenizer.cls_token_id is None or tokenizer.sep_token_id is None:
        raise ValueError("expected DistilBERT CLS/SEP tokenizer")
    raw = tokenizer(text, add_special_tokens=False, truncation=False, return_offsets_mapping=True, verbose=False)
    ids, offsets = raw["input_ids"], raw["offset_mapping"]
    windows = {"input_ids": [], "attention_mask": [], "offset_mapping": []}
    start = 0
    while True:
        end = min(start + capacity, len(ids))
        window = [tokenizer.cls_token_id] + ids[start:end] + [tokenizer.sep_token_id]
        windows["input_ids"].append(window)
        windows["attention_mask"].append([1] * len(window))
        windows["offset_mapping"].append([(0, 0)] + offsets[start:end] + [(0, 0)])
        if end == len(ids):
            break
        start += capacity - stride
    return windows


def decode_window(logits, offsets, text, max_length, has_previous=False, has_next=False):
    scores = np.asarray(logits, dtype=np.float64)
    scores = np.exp(scores - scores.max(axis=-1, keepdims=True))
    scores /= scores.sum(axis=-1, keepdims=True)
    ids = scores.argmax(axis=-1)
    bp = byte_positions(text)
    findings = []
    current = None
    for i, (start, end) in enumerate(offsets):
        value = int(ids[i])
        if start == end or value == 0:
            if current:
                findings.append(current)
                current = None
            continue
        label = LABELS[(value - 1) // 2]
        score = float(scores[i, value])
        if value % 2 == 1 or current is None or current["label"] != label:
            if current:
                findings.append(current)
            current = {"label": label, "start": bp[start], "end": bp[end], "raw_score": score, "confidence": None}
        else:
            current["end"] = bp[end]
            current["raw_score"] = min(current["raw_score"], score)
    if current:
        findings.append(current)
    content = [(a, b) for a, b in offsets if a != b]
    if len(offsets) == max_length and content:
        if any(f["start"] == bp[content[0][0]] and f["end"] == bp[content[-1][1]] for f in findings):
            raise ValueError("entity spans a full window; complete entity boundary is unsupported")
    if content:
        for f in findings:
            f["_boundary_fragment"] = bool((has_previous and f["start"] == bp[content[0][0]]) or (has_next and f["end"] == bp[content[-1][1]]))
    return findings


def resolve_findings(findings):
    # Prefer an observed containing prediction over a clipped window-edge fragment.
    # Do not discard an edge prediction without a same-label containing candidate.
    findings = [f for f in findings if not (f.get("_boundary_fragment", False) and any(
        g["label"] == f["label"] and g["start"] <= f["start"] and g["end"] >= f["end"]
        and (g["start"] < f["start"] or g["end"] > f["end"]) for g in findings))]
    accepted = []
    for finding in sorted(findings, key=lambda f: (-f["raw_score"], f["start"], f["end"], f["label"])):
        if not any(finding["start"] < f["end"] and f["start"] < finding["end"] for f in accepted):
            accepted.append({k:v for k,v in finding.items() if not k.startswith("_")})
    return sorted(accepted, key=lambda f: (f["start"], f["end"], f["label"]))


def result(text, findings, calibration=None, threshold=None):
    findings = resolve_findings(findings)
    if calibration and calibration.get("method") == "isotonic":
        for f in findings:
            f["confidence"] = float(np.interp(f["raw_score"], calibration["xs"], calibration["ys"]))
    if threshold is not None:
        findings = [f for f in findings if (f["confidence"] if f["confidence"] is not None else f["raw_score"]) >= threshold]
    data = text.encode("utf-8")
    parts, cursor = [], 0
    for f in findings:
        parts.extend([data[cursor:f["start"]], b"[REDACTED]"])
        cursor = f["end"]
    parts.append(data[cursor:])
    return {"findings": findings, "contains_pii": bool(findings), "redacted": b"".join(parts).decode("utf-8")}


def load_jsonl(path):
    with open(path) as handle:
        return [json.loads(line) for line in handle if line.strip()]


class Predictor:
    def __init__(self, bundle, device="cpu", backend="onnx"):
        from transformers import AutoTokenizer, AutoModelForTokenClassification
        self.bundle = Path(bundle)
        self.config = json.loads((self.bundle / "config.json").read_text())
        self.tokenizer = AutoTokenizer.from_pretrained(self.bundle, local_files_only=True)
        self.backend, self.device = backend, device
        self.calibration = json.loads((self.bundle / self.config["calibration_file"]).read_text()) if self.config.get("calibration_file") else None
        if backend == "onnx":
            import onnxruntime as ort
            options = ort.SessionOptions()
            options.intra_op_num_threads = 1
            self.session = ort.InferenceSession(str(self.bundle / "model.onnx"), sess_options=options, providers=["CPUExecutionProvider"])
        else:
            self.model = AutoModelForTokenClassification.from_pretrained(self.bundle / "checkpoint", local_files_only=True, attn_implementation="eager").to(device).eval()

    def predict(self, text, calibrated=True):
        encoded = encode_windows(self.tokenizer, text, self.config["max_length"], self.config["stride"])
        findings = []
        for index, (ids, mask, offsets) in enumerate(zip(encoded["input_ids"], encoded["attention_mask"], encoded["offset_mapping"])):
            if self.backend == "onnx":
                logits = self.session.run(None, {"input_ids": np.array([ids], dtype=np.int64), "attention_mask": np.array([mask], dtype=np.int64)})[0][0]
            else:
                import torch
                with torch.no_grad():
                    logits = self.model(input_ids=torch.tensor([ids], device=self.device), attention_mask=torch.tensor([mask], device=self.device)).logits[0].cpu().numpy()
            findings.extend(decode_window(logits, offsets, text, self.config["max_length"], index > 0, index + 1 < len(encoded["input_ids"])))
        calibration = self.calibration if calibrated else None
        return result(text, findings, calibration, calibration.get("threshold") if calibration else None)