app_v2.py

"""
PII Reveal - Document Privacy Explorer (v2)
============================================
Redesigned frontend addressing ui-critique-1.txt:
 - calmer palette, one brand accent, category colors desaturated
 - KPI summary cards (with risk level)
 - tinted category chips + stacked distribution bar
 - premium document viewer with Original / Masked toolbar + focus mode
 - inspection-rail sidebar (Filters -> Findings -> Actions)
 - hover-linked span <-> sidebar inspection
 - unified 8/12/16/24/32 spacing scale
 - Inter typography, polished hierarchy

Backend (model, server, endpoints) is identical to app.py.
"""

# ── stdlib ───────────────────────────────────────────────────────
import dataclasses
import functools
import json
import math
import os
import re
import tempfile
from bisect import bisect_left, bisect_right
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Final

# ── third-party ──────────────────────────────────────────────────
import gradio as gr
import spaces
import tiktoken
import torch
import torch.nn.functional as F
from fastapi import UploadFile, File
from fastapi.responses import HTMLResponse, JSONResponse
from huggingface_hub import snapshot_download
from safetensors import safe_open

# ── configuration ────────────────────────────────────────────────
MODEL_REPO = os.getenv("MODEL_ID", "charles-first-org/second-model")
HF_TOKEN = os.getenv("HF_TOKEN", None)
MODEL_DIR = Path(snapshot_download(MODEL_REPO, token=HF_TOKEN))

# Desaturated category palette (~25% lower saturation than v1), paired with
# tint/text tokens so chips, dots, and span highlights stay visually quiet.
CATEGORIES_META = {
    "private_person":  {"color": "#dc2626", "tint": "rgba(220,38,38,0.08)",  "text": "#991b1b", "label": "Person"},
    "private_address": {"color": "#0891b2", "tint": "rgba(8,145,178,0.08)",  "text": "#155e75", "label": "Address"},
    "private_email":   {"color": "#2563eb", "tint": "rgba(37,99,235,0.08)",  "text": "#1e40af", "label": "Email"},
    "private_phone":   {"color": "#16a34a", "tint": "rgba(22,163,74,0.08)",  "text": "#14532d", "label": "Phone"},
    "private_url":     {"color": "#ca8a04", "tint": "rgba(202,138,4,0.10)",  "text": "#713f12", "label": "URL"},
    "private_date":    {"color": "#9333ea", "tint": "rgba(147,51,234,0.08)", "text": "#6b21a8", "label": "Date"},
    "account_number":  {"color": "#ea580c", "tint": "rgba(234,88,12,0.08)",  "text": "#7c2d12", "label": "Account"},
    "secret":          {"color": "#b91c1c", "tint": "rgba(185,28,28,0.10)",  "text": "#7f1d1d", "label": "Secret"},
}

# =====================================================================
# MODEL  ARCHITECTURE  +  INFERENCE  (from reference implementation)
# =====================================================================

PRIVACY_FILTER_MODEL_TYPE: Final[str] = "privacy_filter"
REQUIRED_MODEL_CONFIG_KEYS: Final[tuple[str, ...]] = (
    "model_type", "encoding", "num_hidden_layers", "num_experts",
    "experts_per_token", "vocab_size", "num_labels", "hidden_size",
    "intermediate_size", "head_dim", "num_attention_heads",
    "num_key_value_heads", "sliding_window", "bidirectional_context",
    "bidirectional_left_context", "bidirectional_right_context",
    "default_n_ctx", "initial_context_length", "rope_theta",
    "rope_scaling_factor", "rope_ntk_alpha", "rope_ntk_beta", "param_dtype",
)
BACKGROUND_CLASS_LABEL: Final[str] = "O"
BOUNDARY_PREFIXES: Final[tuple[str, ...]] = ("B", "I", "E", "S")
SPAN_CLASS_NAMES: Final[tuple[str, ...]] = (
    BACKGROUND_CLASS_LABEL,
    "account_number", "private_address", "private_date", "private_email",
    "private_person", "private_phone", "private_url", "secret",
)
NER_CLASS_NAMES: Final[tuple[str, ...]] = (BACKGROUND_CLASS_LABEL,) + tuple(
    f"{prefix}-{base}"
    for base in SPAN_CLASS_NAMES if base != BACKGROUND_CLASS_LABEL
    for prefix in BOUNDARY_PREFIXES
)
VITERBI_TRANSITION_BIAS_KEYS: Final[tuple[str, ...]] = (
    "transition_bias_background_stay", "transition_bias_background_to_start",
    "transition_bias_inside_to_continue", "transition_bias_inside_to_end",
    "transition_bias_end_to_background", "transition_bias_end_to_start",
)
DEFAULT_VITERBI_CALIBRATION_PRESET: Final[str] = "default"


def validate_model_config_contract(cfg: dict, *, context: str) -> None:
    missing = [k for k in REQUIRED_MODEL_CONFIG_KEYS if k not in cfg]
    if missing:
        raise ValueError(f"{context} missing keys: {', '.join(missing)}")
    if cfg.get("model_type") != PRIVACY_FILTER_MODEL_TYPE:
        raise ValueError(f"{context} model_type must be {PRIVACY_FILTER_MODEL_TYPE!r}")
    if cfg.get("bidirectional_context") is not True:
        raise ValueError(f"{context} must use bidirectional_context=true")
    lc, rc = cfg.get("bidirectional_left_context"), cfg.get("bidirectional_right_context")
    if not isinstance(lc, int) or not isinstance(rc, int) or lc != rc or lc < 0:
        raise ValueError(f"{context} bidirectional context must be equal non-negative ints")
    sw = cfg.get("sliding_window")
    if sw != 2 * lc + 1:
        raise ValueError(f"{context} sliding_window must equal 2*context+1")
    if cfg["num_labels"] != 33:
        raise ValueError(f"{context} num_labels must be 33")
    if cfg["param_dtype"] != "bfloat16":
        raise ValueError(f"{context} param_dtype must be bfloat16")


# ── model helpers ────────────────────────────────────────────────

def expert_linear(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor:
    n, e, k = x.shape
    _, _, _, o = weight.shape
    out = torch.bmm(x.reshape(n * e, 1, k), weight.reshape(n * e, k, o)).reshape(n, e, o)
    return out + bias if bias is not None else out


@dataclass
class ModelConfig:
    num_hidden_layers: int; num_experts: int; experts_per_token: int
    vocab_size: int; num_labels: int; hidden_size: int; intermediate_size: int
    head_dim: int; num_attention_heads: int; num_key_value_heads: int
    bidirectional_context_size: int; initial_context_length: int
    rope_theta: float; rope_scaling_factor: float; rope_ntk_alpha: float; rope_ntk_beta: float

    @classmethod
    def from_checkpoint_config(cls, cfg: dict, *, context: str) -> "ModelConfig":
        cfg = dict(cfg)
        cfg["bidirectional_context_size"] = cfg["bidirectional_left_context"]
        fields = {f.name for f in dataclasses.fields(cls)}
        return cls(**{k: v for k, v in cfg.items() if k in fields})


class RMSNorm(torch.nn.Module):
    def __init__(self, n: int, eps: float = 1e-5, device=None):
        super().__init__()
        self.eps = eps
        self.scale = torch.nn.Parameter(torch.ones(n, device=device, dtype=torch.float32))

    def forward(self, x):
        t = x.float()
        return (t * torch.rsqrt(t.pow(2).mean(-1, keepdim=True) + self.eps) * self.scale).to(x.dtype)


def apply_rope(x, cos, sin):
    cos = cos.unsqueeze(-2).to(x.dtype); sin = sin.unsqueeze(-2).to(x.dtype)
    x1, x2 = x[..., ::2], x[..., 1::2]
    return torch.stack((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1).reshape(x.shape)


class RotaryEmbedding(torch.nn.Module):
    def __init__(self, head_dim, base, dtype, *, initial_context_length=4096,
                 scaling_factor=1.0, ntk_alpha=1.0, ntk_beta=32.0, device=None):
        super().__init__()
        self.head_dim, self.base, self.dtype = head_dim, base, dtype
        self.initial_context_length = initial_context_length
        self.scaling_factor, self.ntk_alpha, self.ntk_beta = scaling_factor, ntk_alpha, ntk_beta
        self.device = device
        mp = max(int(initial_context_length * scaling_factor), initial_context_length)
        self.max_position_embeddings = mp
        cos, sin = self._compute(mp, device=torch.device("cpu"))
        target = device or torch.device("cpu")
        self.register_buffer("cos_cache", cos.to(target), persistent=False)
        self.register_buffer("sin_cache", sin.to(target), persistent=False)

    def _inv_freq(self, device=None):
        device = device or self.device
        freq = self.base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float, device=device) / self.head_dim)
        if self.scaling_factor > 1.0:
            d_half = self.head_dim / 2
            low = d_half * math.log(self.initial_context_length / (self.ntk_beta * 2 * math.pi)) / math.log(self.base)
            high = d_half * math.log(self.initial_context_length / (self.ntk_alpha * 2 * math.pi)) / math.log(self.base)
            interp = 1.0 / (self.scaling_factor * freq)
            extrap = 1.0 / freq
            ramp = (torch.arange(d_half, dtype=torch.float32, device=device) - low) / (high - low)
            mask = 1 - ramp.clamp(0, 1)
            return interp * (1 - mask) + extrap * mask
        return 1.0 / freq

    def _compute(self, n, device=None):
        inv_freq = self._inv_freq(device)
        t = torch.arange(n, dtype=torch.float32, device=device or self.device)
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        c = 0.1 * math.log(self.scaling_factor) + 1.0 if self.scaling_factor > 1.0 else 1.0
        return (freqs.cos() * c).to(self.dtype), (freqs.sin() * c).to(self.dtype)

    def forward(self, q, k):
        n = q.shape[0]
        if n > self.cos_cache.shape[0]:
            cos, sin = self._compute(n, torch.device("cpu"))
            self.cos_cache, self.sin_cache = cos.to(q.device), sin.to(q.device)
        cc = self.cos_cache.to(q.device) if self.cos_cache.device != q.device else self.cos_cache
        sc = self.sin_cache.to(q.device) if self.sin_cache.device != q.device else self.sin_cache
        cos, sin = cc[:n], sc[:n]
        q = apply_rope(q.view(n, -1, self.head_dim), cos, sin).reshape(q.shape)
        k = apply_rope(k.view(n, -1, self.head_dim), cos, sin).reshape(k.shape)
        return q, k


def sdpa(Q, K, V, S, sm_scale, ctx):
    n, nh, qm, hd = Q.shape
    w = 2 * ctx + 1
    Kp = F.pad(K, (0, 0, 0, 0, ctx, ctx)); Vp = F.pad(V, (0, 0, 0, 0, ctx, ctx))
    Kw = Kp.unfold(0, w, 1).permute(0, 3, 1, 2); Vw = Vp.unfold(0, w, 1).permute(0, 3, 1, 2)
    idx = torch.arange(w, device=Q.device) - ctx
    pos = torch.arange(n, device=Q.device)[:, None] + idx[None, :]
    valid = (pos >= 0) & (pos < n)
    scores = torch.einsum("nhqd,nwhd->nhqw", Q, Kw).float() * sm_scale
    scores = scores.masked_fill(~valid[:, None, None, :], -float("inf"))
    sink = (S * math.log(2.0)).reshape(nh, qm)[None, :, :, None].expand(n, -1, -1, 1)
    scores = torch.cat([scores, sink], dim=-1)
    wt = torch.softmax(scores, dim=-1)[..., :-1].to(V.dtype)
    return torch.einsum("nhqw,nwhd->nhqd", wt, Vw).reshape(n, -1)


class AttentionBlock(torch.nn.Module):
    def __init__(self, cfg: ModelConfig, device=None):
        super().__init__()
        dt = torch.bfloat16
        self.head_dim, self.nah, self.nkv = cfg.head_dim, cfg.num_attention_heads, cfg.num_key_value_heads
        self.ctx = int(cfg.bidirectional_context_size)
        self.sinks = torch.nn.Parameter(torch.empty(cfg.num_attention_heads, device=device, dtype=torch.float32))
        self.norm = RMSNorm(cfg.hidden_size, device=device)
        qkv_d = cfg.head_dim * (cfg.num_attention_heads + 2 * cfg.num_key_value_heads)
        self.qkv = torch.nn.Linear(cfg.hidden_size, qkv_d, device=device, dtype=dt)
        self.out = torch.nn.Linear(cfg.head_dim * cfg.num_attention_heads, cfg.hidden_size, device=device, dtype=dt)
        self.qk_scale = 1 / math.sqrt(math.sqrt(cfg.head_dim))
        self.rope = RotaryEmbedding(cfg.head_dim, int(cfg.rope_theta), torch.float32,
                                     initial_context_length=cfg.initial_context_length,
                                     scaling_factor=cfg.rope_scaling_factor,
                                     ntk_alpha=cfg.rope_ntk_alpha, ntk_beta=cfg.rope_ntk_beta, device=device)

    def forward(self, x):
        t = self.norm(x).to(self.qkv.weight.dtype)
        qkv = F.linear(t, self.qkv.weight, self.qkv.bias)
        hd, nah, nkv = self.head_dim, self.nah, self.nkv
        q = qkv[:, :nah * hd].contiguous()
        k = qkv[:, nah * hd:(nah + nkv) * hd].contiguous()
        v = qkv[:, (nah + nkv) * hd:(nah + 2 * nkv) * hd].contiguous()
        q, k = self.rope(q, k)
        q, k = q * self.qk_scale, k * self.qk_scale
        n = q.shape[0]
        q = q.view(n, nkv, nah // nkv, hd); k = k.view(n, nkv, hd); v = v.view(n, nkv, hd)
        ao = sdpa(q, k, v, self.sinks, 1.0, self.ctx).to(self.out.weight.dtype)
        return x + F.linear(ao, self.out.weight, self.out.bias).to(x.dtype)


def swiglu(x, alpha=1.702, limit=7.0):
    g, l = x.chunk(2, dim=-1)
    g, l = g.clamp(max=limit), l.clamp(-limit, limit)
    return g * torch.sigmoid(alpha * g) * (l + 1)


class MLPBlock(torch.nn.Module):
    def __init__(self, cfg: ModelConfig, device=None):
        super().__init__()
        dt = torch.bfloat16
        self.ne, self.ept = cfg.num_experts, cfg.experts_per_token
        self.norm = RMSNorm(cfg.hidden_size, device=device)
        self.gate = torch.nn.Linear(cfg.hidden_size, cfg.num_experts, device=device, dtype=dt)
        self.mlp1_weight = torch.nn.Parameter(torch.empty(cfg.num_experts, cfg.hidden_size, cfg.intermediate_size * 2, device=device, dtype=dt))
        self.mlp1_bias = torch.nn.Parameter(torch.empty(cfg.num_experts, cfg.intermediate_size * 2, device=device, dtype=dt))
        self.mlp2_weight = torch.nn.Parameter(torch.empty(cfg.num_experts, cfg.intermediate_size, cfg.hidden_size, device=device, dtype=dt))
        self.mlp2_bias = torch.nn.Parameter(torch.empty(cfg.num_experts, cfg.hidden_size, device=device, dtype=dt))

    def forward(self, x):
        t = self.norm(x)
        gs = F.linear(t.float(), self.gate.weight.float(), self.gate.bias.float())
        top = torch.topk(gs, k=self.ept, dim=-1, sorted=True)
        ew = torch.softmax(top.values, dim=-1) / self.ept
        ei = top.indices
        ept = self.ept

        def _chunk(tc, eic, ewc):
            o = expert_linear(tc.float().unsqueeze(1).expand(-1, eic.shape[1], -1),
                              self.mlp1_weight[eic].float(), self.mlp1_bias[eic].float())
            o = swiglu(o)
            o = expert_linear(o.float(), self.mlp2_weight[eic].float(), self.mlp2_bias[eic].float())
            return (torch.einsum("bec,be->bc", o.to(ewc.dtype), ewc) * ept).to(x.dtype)

        cs = 32
        if t.shape[0] > cs:
            parts = [_chunk(t[s:s+cs], ei[s:s+cs], ew[s:s+cs]) for s in range(0, t.shape[0], cs)]
            return x + torch.cat(parts, 0)
        return x + _chunk(t, ei, ew)


class TransformerBlock(torch.nn.Module):
    def __init__(self, cfg, device=None):
        super().__init__()
        self.attn = AttentionBlock(cfg, device=device)
        self.mlp = MLPBlock(cfg, device=device)
    def forward(self, x):
        return self.mlp(self.attn(x))


class Checkpoint:
    @staticmethod
    def build_param_name_map(n):
        return ({f"block.{i}.mlp.mlp1_bias": f"block.{i}.mlp.swiglu.bias" for i in range(n)}
              | {f"block.{i}.mlp.mlp1_weight": f"block.{i}.mlp.swiglu.weight" for i in range(n)}
              | {f"block.{i}.mlp.mlp2_bias": f"block.{i}.mlp.out.bias" for i in range(n)}
              | {f"block.{i}.mlp.mlp2_weight": f"block.{i}.mlp.out.weight" for i in range(n)})

    def __init__(self, path, device, num_hidden_layers):
        self.pnm = self.build_param_name_map(num_hidden_layers)
        self.ds = device.type if device.index is None else f"{device.type}:{device.index}"
        files = [os.path.join(path, f) for f in os.listdir(path) if f.endswith(".safetensors")]
        self.map = {}
        for sf in files:
            with safe_open(sf, framework="pt", device=self.ds) as h:
                for k in h.keys():
                    self.map[k] = sf

    def get(self, name):
        mapped = self.pnm.get(name, name)
        with safe_open(self.map[mapped], framework="pt", device=self.ds) as h:
            return h.get_tensor(mapped)


class Transformer(torch.nn.Module):
    def __init__(self, cfg, device):
        super().__init__()
        dt = torch.bfloat16
        self.embedding = torch.nn.Embedding(cfg.vocab_size, cfg.hidden_size, device=device, dtype=dt)
        self.block = torch.nn.ModuleList([TransformerBlock(cfg, device=device) for _ in range(cfg.num_hidden_layers)])
        self.norm = RMSNorm(cfg.hidden_size, device=device)
        self.unembedding = torch.nn.Linear(cfg.hidden_size, cfg.num_labels, bias=False, device=device, dtype=dt)

    def forward(self, token_ids):
        x = self.embedding(token_ids)
        for blk in self.block:
            x = blk(x)
        return F.linear(self.norm(x), self.unembedding.weight, None)

    @classmethod
    def from_checkpoint(cls, checkpoint_dir, *, device):
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
        torch.set_float32_matmul_precision("highest")
        cp = json.loads((Path(checkpoint_dir) / "config.json").read_text())
        validate_model_config_contract(cp, context=str(checkpoint_dir))
        cfg = ModelConfig.from_checkpoint_config(cp, context=str(checkpoint_dir))
        ckpt = Checkpoint(checkpoint_dir, device, cfg.num_hidden_layers)
        m = cls(cfg, device); m.eval()
        for name, param in m.named_parameters():
            loaded = ckpt.get(name)
            if param.shape != loaded.shape:
                raise ValueError(f"Shape mismatch {name}: {param.shape} vs {loaded.shape}")
            param.data.copy_(loaded)
        return m


# ── label info + span decoding ───────────────────────────────────

@dataclass(frozen=True)
class LabelInfo:
    boundary_label_lookup: dict[str, dict[str, int]]
    token_to_span_label: dict[int, int]
    token_boundary_tags: dict[int, str | None]
    span_class_names: tuple[str, ...]
    span_label_lookup: dict[str, int]
    background_token_label: int
    background_span_label: int


def labels_to_spans(labels_by_index, label_info):
    spans, cur_label, start_idx, prev_idx = [], None, None, None
    bg = label_info.background_span_label
    for ti in sorted(labels_by_index):
        lid = labels_by_index[ti]
        sl = label_info.token_to_span_label.get(lid)
        bt = label_info.token_boundary_tags.get(lid)
        if prev_idx is not None and ti != prev_idx + 1:
            if cur_label is not None and start_idx is not None:
                spans.append((cur_label, start_idx, prev_idx + 1))
            cur_label = start_idx = None
        if sl is None:
            prev_idx = ti; continue
        if sl == bg:
            if cur_label is not None and start_idx is not None:
                spans.append((cur_label, start_idx, ti))
            cur_label = start_idx = None; prev_idx = ti; continue
        if bt == "S":
            if cur_label is not None and start_idx is not None and prev_idx is not None:
                spans.append((cur_label, start_idx, prev_idx + 1))
            spans.append((sl, ti, ti + 1)); cur_label = start_idx = None
        elif bt == "B":
            if cur_label is not None and start_idx is not None and prev_idx is not None:
                spans.append((cur_label, start_idx, prev_idx + 1))
            cur_label, start_idx = sl, ti
        elif bt == "I":
            if cur_label is None or cur_label != sl:
                if cur_label is not None and start_idx is not None and prev_idx is not None:
                    spans.append((cur_label, start_idx, prev_idx + 1))
                cur_label, start_idx = sl, ti
        elif bt == "E":
            if cur_label is None or cur_label != sl or start_idx is None:
                if cur_label is not None and start_idx is not None and prev_idx is not None:
                    spans.append((cur_label, start_idx, prev_idx + 1))
                spans.append((sl, ti, ti + 1)); cur_label = start_idx = None
            else:
                spans.append((cur_label, start_idx, ti + 1)); cur_label = start_idx = None
        else:
            if cur_label is not None and start_idx is not None and prev_idx is not None:
                spans.append((cur_label, start_idx, prev_idx + 1))
            cur_label = start_idx = None
        prev_idx = ti
    if cur_label is not None and start_idx is not None and prev_idx is not None:
        spans.append((cur_label, start_idx, prev_idx + 1))
    return spans


def token_spans_to_char_spans(spans, cs, ce):
    out = []
    for li, ts, te in spans:
        if not (0 <= ts < te <= len(cs)):
            continue
        s, e = cs[ts], ce[te - 1]
        if e > s:
            out.append((li, s, e))
    return out


def trim_char_spans_whitespace(spans, text):
    out = []
    for li, s, e in spans:
        if not (0 <= s < e <= len(text)):
            continue
        while s < e and text[s].isspace(): s += 1
        while e > s and text[e - 1].isspace(): e -= 1
        if e > s:
            out.append((li, s, e))
    return out


# ── viterbi decoder ──────────────────────────────────────────────

@functools.lru_cache(maxsize=1)
def get_viterbi_transition_biases():
    cp = MODEL_DIR / "viterbi_calibration.json"
    default = {k: 0.0 for k in VITERBI_TRANSITION_BIAS_KEYS}
    if not cp.is_file():
        return default
    payload = json.loads(cp.read_text())
    raw = payload
    ops = payload.get("operating_points")
    if isinstance(ops, dict):
        preset = ops.get(DEFAULT_VITERBI_CALIBRATION_PRESET)
        if isinstance(preset, dict):
            raw = preset.get("biases", raw)
    if not isinstance(raw, dict):
        return default
    return {k: float(raw.get(k, 0.0)) for k in VITERBI_TRANSITION_BIAS_KEYS}


class Decoder:
    def __init__(self, label_info):
        nc = len(label_info.token_to_span_label)
        self._start = torch.full((nc,), -1e9, dtype=torch.float32)
        self._end = torch.full((nc,), -1e9, dtype=torch.float32)
        self._trans = torch.full((nc, nc), -1e9, dtype=torch.float32)
        biases = get_viterbi_transition_biases()
        bg_tok, bg_sp = label_info.background_token_label, label_info.background_span_label
        ttsl, tbt = label_info.token_to_span_label, label_info.token_boundary_tags
        for i in range(nc):
            tag, sl = tbt.get(i), ttsl.get(i)
            if tag in {"B", "S"} or i == bg_tok: self._start[i] = 0.0
            if tag in {"E", "S"} or i == bg_tok: self._end[i] = 0.0
            for j in range(nc):
                nt, ns = tbt.get(j), ttsl.get(j)
                if self._valid(tag, sl, nt, ns, bg_tok, bg_sp, j):
                    self._trans[i, j] = self._bias(tag, sl, nt, ns, bg_sp, biases)

    @staticmethod
    def _valid(pt, ps, nt, ns, bti, bsi, ni):
        nb = ns == bsi or ni == bti
        if (ns is None or nt is None) and not nb: return False
        if pt is None or ps is None: return nb or nt in {"B", "S"}
        if ps == bsi or pt in {"E", "S"}: return nb or nt in {"B", "S"}
        if pt in {"B", "I"}: return ps == ns and nt in {"I", "E"}
        return False

    @staticmethod
    def _bias(pt, ps, nt, ns, bsi, b):
        nb, pb = ns == bsi, ps == bsi
        if pb: return b["transition_bias_background_stay"] if nb else b["transition_bias_background_to_start"]
        if pt in {"B", "I"}: return b["transition_bias_inside_to_continue"] if nt == "I" else b["transition_bias_inside_to_end"]
        return b["transition_bias_end_to_background"] if nb else b["transition_bias_end_to_start"]

    def decode(self, lp):
        sl, nc = lp.shape
        if sl == 0: return []
        st = self._start.to(lp.device, lp.dtype)
        en = self._end.to(lp.device, lp.dtype)
        tr = self._trans.to(lp.device, lp.dtype)
        scores = lp[0] + st
        bp = torch.empty((sl - 1, nc), device=lp.device, dtype=torch.int64)
        for i in range(1, sl):
            t = scores.unsqueeze(1) + tr
            bs, bi = t.max(dim=0)
            scores = bs + lp[i]; bp[i - 1] = bi
        if not torch.isfinite(scores).any(): return lp.argmax(dim=1).tolist()
        scores += en
        path = torch.empty(sl, device=lp.device, dtype=torch.int64)
        path[-1] = scores.argmax()
        for i in range(sl - 2, -1, -1): path[i] = bp[i, path[i + 1]]
        return path.tolist()


# ── runtime singleton ────────────────────────────────────────────

@dataclass(frozen=True)
class InferenceRuntime:
    model: Transformer; encoding: tiktoken.Encoding; label_info: LabelInfo
    device: torch.device; n_ctx: int


@functools.lru_cache(maxsize=1)
def get_runtime():
    cp = MODEL_DIR
    cfg = json.loads((cp / "config.json").read_text())
    validate_model_config_contract(cfg, context=str(cp))
    device = torch.device("cuda")
    encoding = tiktoken.get_encoding(str(cfg["encoding"]).strip())
    scn = [BACKGROUND_CLASS_LABEL]; sll = {BACKGROUND_CLASS_LABEL: 0}
    bll, ttsl, tbt = {}, {}, {}
    bg_idx = None
    for idx, name in enumerate(NER_CLASS_NAMES):
        if name == BACKGROUND_CLASS_LABEL:
            bg_idx = idx; ttsl[idx] = 0; tbt[idx] = None; continue
        bnd, base = name.split("-", 1)
        si = sll.get(base)
        if si is None:
            si = len(scn); scn.append(base); sll[base] = si
        ttsl[idx] = si; tbt[idx] = bnd
        bll.setdefault(base, {})[bnd] = idx
    li = LabelInfo(bll, ttsl, tbt, tuple(scn), sll, bg_idx, 0)
    m = Transformer.from_checkpoint(str(cp), device=device)
    return InferenceRuntime(m, encoding, li, device, int(cfg["default_n_ctx"]))


@torch.inference_mode()
def predict_text(runtime, text, decoder):
    tids = tuple(int(t) for t in runtime.encoding.encode(text, allowed_special="all"))
    if not tids: return text, []
    scores = []
    for s in range(0, len(tids), runtime.n_ctx):
        e = min(s + runtime.n_ctx, len(tids))
        wt = torch.tensor(tids[s:e], device=runtime.device, dtype=torch.int32)
        lp = F.log_softmax(runtime.model(wt).float(), dim=-1)
        scores.extend(lp.unbind(0))
    stacked = torch.stack(scores, 0)
    dl = decoder.decode(stacked)
    if len(dl) != len(tids): dl = stacked.argmax(dim=1).tolist()
    pli = {i: int(l) for i, l in enumerate(dl)}
    pts = labels_to_spans(pli, runtime.label_info)
    tb = [runtime.encoding.decode_single_token_bytes(t) for t in tids]
    dt = b"".join(tb).decode("utf-8", errors="replace")
    cbs, cbe = [], []
    bc = 0
    for ch in dt: cbs.append(bc); bc += len(ch.encode("utf-8")); cbe.append(bc)
    cs, ce = [], []
    tbc = 0
    for rb in tb:
        tbs = tbc; tbe = tbs + len(rb); tbc = tbe
        cs.append(bisect_right(cbe, tbs)); ce.append(bisect_left(cbs, tbe))
    pcs = token_spans_to_char_spans(pts, cs, ce)
    pcs = trim_char_spans_whitespace(pcs, dt if dt != text else text)
    src = dt if dt != text else text
    detected = []
    for li, s, e in pcs:
        if 0 <= li < len(runtime.label_info.span_class_names):
            lbl = runtime.label_info.span_class_names[li]
        else:
            lbl = f"label_{li}"
        detected.append({"label": lbl, "start": s, "end": e, "text": src[s:e]})
    return src, detected


# =====================================================================
# APPLICATION  LAYER
# =====================================================================

def extract_text(file_path: str) -> str:
    suffix = Path(file_path).suffix.lower()
    if suffix == ".pdf":
        import fitz
        doc = fitz.open(file_path)
        pages = [page.get_text() for page in doc]
        doc.close()
        return "\n\n".join(pages)
    elif suffix in (".docx", ".doc"):
        from docx import Document
        doc = Document(file_path)
        return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip())
    raise ValueError(f"Unsupported file type: {suffix}")


def compute_stats(text, spans):
    total = len(text)
    pii_chars = sum(s["end"] - s["start"] for s in spans)
    by_cat = {}
    for s in spans:
        c = s["label"]
        by_cat.setdefault(c, {"count": 0, "chars": 0})
        by_cat[c]["count"] += 1; by_cat[c]["chars"] += s["end"] - s["start"]
    pct = round(pii_chars / total * 100, 1) if total else 0
    # Risk tiering — secrets/accounts/emails make a document higher-risk even
    # at low coverage, so combine a density rule with a sensitive-category rule.
    sensitive = sum(by_cat.get(k, {}).get("count", 0) for k in ("secret", "account_number", "private_email"))
    if pct >= 15 or sensitive >= 5:
        risk = "High"
    elif pct >= 5 or sensitive >= 1 or len(spans) >= 10:
        risk = "Medium"
    else:
        risk = "Low"
    return {
        "total_chars": total, "pii_chars": pii_chars,
        "pii_percentage": pct,
        "total_spans": len(spans), "categories": by_cat, "num_categories": len(by_cat),
        "risk_level": risk,
    }


def detect_speakers(text, spans):
    patterns = [r"^([A-Z][a-zA-Z ]{1,30}):\s", r"^\[([^\]]{1,30})\]\s", r"^(Speaker\s*\d+):\s"]
    line_sp, pos, cur = [], 0, None
    for line in text.split("\n"):
        for p in patterns:
            m = re.match(p, line)
            if m: cur = m.group(1).strip(); break
        line_sp.append((pos, pos + len(line), cur)); pos += len(line) + 1
    result = {}
    for span in spans:
        mid = (span["start"] + span["end"]) // 2
        speaker = "Document"
        for ls, le, sp in line_sp:
            if ls <= mid <= le and sp: speaker = sp; break
        result[speaker] = result.get(speaker, 0) + 1
    return {} if list(result.keys()) == ["Document"] else result


@spaces.GPU
def run_pii_analysis(text: str):
    """GPU-accelerated PII detection."""
    runtime = get_runtime()
    decoder = Decoder(label_info=runtime.label_info)
    source_text, detected = predict_text(runtime, text, decoder)
    return source_text, detected


# ── Gradio Server ────────────────────────────────────────────────
server = gr.Server()


@server.get("/", response_class=HTMLResponse)
async def homepage():
    return FRONTEND_HTML


@server.post("/api/analyze")
async def analyze_document(file: UploadFile = File(...)):
    suffix = Path(file.filename).suffix.lower()
    if suffix not in (".pdf", ".doc", ".docx"):
        return JSONResponse({"error": f"Unsupported: {suffix}. Use PDF, DOC, or DOCX."}, 400)
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp.write(await file.read()); tmp_path = tmp.name
    try:
        text = extract_text(tmp_path)
        if not text.strip():
            return JSONResponse({"error": "No text content found."}, 400)
        source_text, spans = run_pii_analysis(text)
        stats = compute_stats(source_text, spans)
        speakers = detect_speakers(source_text, spans)
        return JSONResponse({
            "filename": file.filename, "text": source_text, "spans": spans,
            "stats": stats, "speakers": speakers,
            "categories_meta": {k: {"color": v["color"], "tint": v["tint"], "text": v["text"], "label": v["label"]}
                                for k, v in CATEGORIES_META.items()},
        })
    except Exception as e:
        return JSONResponse({"error": str(e)}, 500)
    finally:
        if os.path.exists(tmp_path): os.unlink(tmp_path)


@server.api(name="analyze_text")
def analyze_text_api(text: str) -> str:
    """Gradio API: analyze raw text for PII."""
    source_text, spans = run_pii_analysis(text)
    stats = compute_stats(source_text, spans)
    return json.dumps({"text": source_text, "spans": spans, "stats": stats}, ensure_ascii=False)


# ── Frontend HTML (redesigned) ───────────────────────────────────
FRONTEND_HTML = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>PII Reveal</title>
<link rel="preconnect" href="https://fonts.googleapis.com">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&family=JetBrains+Mono:wght@400;500&display=swap" rel="stylesheet">
<style>
*,*::before,*::after{box-sizing:border-box;margin:0;padding:0}

:root{
  /* Neutral base */
  --bg:           #f7f7f9;
  --surface:      #ffffff;
  --surface-2:    #fafbfc;
  --surface-warm: #fdfdfb;
  --border:       #e4e7ec;
  --border-soft:  #eef0f3;
  --text:         #0f172a;
  --text-2:       #475569;
  --text-3:       #94a3b8;

  /* Brand accent — one gradient, used sparingly */
  --brand:        #7c3aed;
  --brand-2:      #ec4899;
  --brand-soft:   rgba(124,58,237,.08);
  --brand-ring:   rgba(124,58,237,.22);

  /* Risk */
  --risk-high:    #b91c1c;
  --risk-med:     #b45309;
  --risk-low:     #15803d;

  /* Spacing scale — 8 / 12 / 16 / 24 / 32 / 48 */
  --s-1: 8px; --s-2: 12px; --s-3: 16px; --s-4: 24px; --s-5: 32px; --s-6: 48px;

  /* Radius */
  --r-sm: 8px; --r-md: 12px; --r-lg: 16px; --r-xl: 20px;

  /* Shadow — subtle */
  --shadow-xs: 0 1px 2px rgba(15,23,42,.04);
  --shadow-sm: 0 1px 3px rgba(15,23,42,.06), 0 1px 2px rgba(15,23,42,.04);
  --shadow-md: 0 4px 16px rgba(15,23,42,.06), 0 2px 4px rgba(15,23,42,.04);
  --shadow-lg: 0 12px 40px rgba(15,23,42,.10);
}

html,body{height:100%}
body{
  font-family:'Inter',system-ui,-apple-system,sans-serif;
  background:var(--bg);
  color:var(--text);
  font-feature-settings:"cv11","ss01","ss03";
  font-size:15px;
  line-height:1.5;
  -webkit-font-smoothing:antialiased;
}
button{font-family:inherit}

/* =============== UPLOAD VIEW =============== */
#upload-view{
  display:flex;align-items:center;justify-content:center;
  min-height:100vh;padding:var(--s-5);
}
.upload-card{
  background:var(--surface);
  border:1px solid var(--border);
  border-radius:var(--r-xl);
  padding:var(--s-6) var(--s-5);
  max-width:600px;width:100%;
  box-shadow:var(--shadow-lg);
  text-align:center;
}
.brand{display:inline-flex;align-items:center;gap:var(--s-2)}
.brand-logo{
  width:36px;height:36px;border-radius:10px;
  background:linear-gradient(135deg,var(--brand) 0%,var(--brand-2) 100%);
  display:flex;align-items:center;justify-content:center;
  color:#fff;font-weight:700;font-size:16px;
  box-shadow:0 4px 12px rgba(124,58,237,.25);
}
.brand-name{
  font-size:18px;font-weight:700;letter-spacing:-.01em;
  background:linear-gradient(135deg,var(--brand),var(--brand-2));
  -webkit-background-clip:text;-webkit-text-fill-color:transparent;
}
.upload-card .brand{margin-bottom:var(--s-2)}
.upload-hero{
  font-size:28px;font-weight:700;letter-spacing:-.02em;
  margin-top:var(--s-2);
}
.upload-sub{color:var(--text-2);margin-top:var(--s-1);font-size:15px}

.dropzone{
  margin-top:var(--s-5);
  border:1.5px dashed var(--border);
  background:var(--surface-2);
  border-radius:var(--r-lg);
  padding:var(--s-5) var(--s-4);
  cursor:pointer;transition:all .18s;
  position:relative;
}
.dropzone:hover,.dropzone.dragover{
  border-color:var(--brand);background:var(--brand-soft);
}
.dropzone-icon{
  width:44px;height:44px;margin:0 auto var(--s-2);
  display:flex;align-items:center;justify-content:center;
  background:var(--surface);border:1px solid var(--border);border-radius:12px;
  color:var(--brand);
}
.dropzone-text{font-weight:600;font-size:15px}
.dropzone-hint{color:var(--text-3);font-size:13px;margin-top:4px}
.dropzone input{position:absolute;inset:0;opacity:0;cursor:pointer}

.features{
  display:grid;grid-template-columns:repeat(3,1fr);
  gap:var(--s-2);margin-top:var(--s-5);text-align:left;
}
.feature{
  background:var(--surface-2);
  border:1px solid var(--border-soft);
  padding:var(--s-2) var(--s-3);
  border-radius:var(--r-sm);
}
.feature-title{font-weight:600;font-size:12px}
.feature-desc{color:var(--text-2);font-size:12px;margin-top:2px;line-height:1.45}
.powered-by{margin-top:var(--s-4);font-size:12px;color:var(--text-3)}
.powered-by strong{color:var(--text-2);font-weight:600}

/* =============== RESULTS VIEW =============== */
#results-view{display:none;min-height:100vh}

/* ----- header ----- */
.topbar{
  background:var(--surface);
  border-bottom:1px solid var(--border);
  padding:var(--s-2) var(--s-4);
  display:flex;align-items:center;gap:var(--s-3);
  position:sticky;top:0;z-index:50;
  min-height:64px;
}
.topbar .brand{gap:var(--s-2)}
.topbar .brand-name{font-size:16px}
.file-pill{
  display:inline-flex;align-items:center;gap:var(--s-1);
  padding:6px var(--s-2);
  background:var(--surface-2);
  border:1px solid var(--border-soft);
  border-radius:999px;
  font-size:13px;color:var(--text-2);font-weight:500;
  max-width:380px;
}
.file-pill svg{flex-shrink:0;color:var(--text-3)}
.file-pill-name{overflow:hidden;text-overflow:ellipsis;white-space:nowrap}
.topbar-summary{font-size:13px;color:var(--text-3);margin-left:var(--s-1)}
.topbar-spacer{flex:1}
.topbar-actions{display:flex;gap:var(--s-1)}

.btn{
  display:inline-flex;align-items:center;gap:6px;
  padding:8px 14px;border-radius:var(--r-sm);
  font-weight:500;font-size:13px;
  border:1px solid transparent;cursor:pointer;
  transition:all .15s;
  background:transparent;color:var(--text);
}
.btn-ghost{background:var(--surface);border-color:var(--border)}
.btn-ghost:hover{background:var(--surface-2);border-color:var(--text-3)}
.btn-primary{
  background:linear-gradient(135deg,var(--brand),var(--brand-2));
  color:#fff;font-weight:600;
  box-shadow:0 2px 8px rgba(124,58,237,.22);
}
.btn-primary:hover{filter:brightness(1.06);box-shadow:0 4px 14px rgba(124,58,237,.28)}
.btn svg{width:14px;height:14px}

/* ----- summary cards ----- */
.metrics{
  display:grid;
  grid-template-columns:repeat(4,minmax(0,1fr));
  gap:var(--s-2);
  padding:var(--s-3) var(--s-4);
  background:var(--bg);
}
.metric{
  background:var(--surface);
  border:1px solid var(--border);
  border-radius:var(--r-md);
  padding:var(--s-3);
  display:flex;flex-direction:column;gap:6px;
  position:relative;overflow:hidden;
}
.metric-label{
  font-size:11px;font-weight:600;
  color:var(--text-3);
  letter-spacing:.06em;text-transform:uppercase;
}
.metric-value{
  font-size:28px;font-weight:700;letter-spacing:-.02em;
  color:var(--text);font-variant-numeric:tabular-nums;
}
.metric-hint{font-size:12px;color:var(--text-3)}
.metric-risk{display:inline-flex;align-items:center;gap:6px}
.metric-risk .dot{width:10px;height:10px;border-radius:50%}
.metric-risk.high .dot{background:var(--risk-high);box-shadow:0 0 0 4px rgba(185,28,28,.12)}
.metric-risk.medium .dot{background:var(--risk-med);box-shadow:0 0 0 4px rgba(180,83,9,.12)}
.metric-risk.low .dot{background:var(--risk-low);box-shadow:0 0 0 4px rgba(21,128,61,.12)}
.metric-risk.high .lvl{color:var(--risk-high)}
.metric-risk.medium .lvl{color:var(--risk-med)}
.metric-risk.low .lvl{color:var(--risk-low)}

/* ----- legend + distribution ----- */
.legend{
  background:var(--surface);
  border:1px solid var(--border);
  border-radius:var(--r-md);
  margin:0 var(--s-4) var(--s-3);
  padding:var(--s-3);
}
.legend-header{
  display:flex;align-items:center;justify-content:space-between;
  margin-bottom:var(--s-2);
}
.section-label{
  font-size:11px;font-weight:600;
  color:var(--text-3);
  letter-spacing:.08em;text-transform:uppercase;
}
.dist-bar{
  display:flex;height:6px;border-radius:999px;overflow:hidden;
  background:var(--border-soft);margin-bottom:var(--s-2);
}
.dist-seg{height:100%;transition:opacity .15s}
.dist-seg:hover{opacity:.85}
.chips{display:flex;flex-wrap:wrap;gap:6px}
.chip{
  display:inline-flex;align-items:center;gap:6px;
  padding:4px 10px;
  border-radius:999px;
  font-size:12px;font-weight:500;
  cursor:pointer;user-select:none;
  border:1px solid transparent;
  transition:all .15s;
}
.chip .chip-dot{width:7px;height:7px;border-radius:50%;flex-shrink:0}
.chip .chip-count{
  font-variant-numeric:tabular-nums;
  font-weight:600;opacity:.7;margin-left:2px;
}
.chip.inactive{opacity:.42;filter:grayscale(.3)}

/* ----- layout ----- */
.layout{
  display:grid;
  grid-template-columns:1fr 320px;
  gap:var(--s-3);
  padding:0 var(--s-4) var(--s-4);
  min-height:calc(100vh - 260px);
}

/* ----- document viewer ----- */
.doc-shell{
  background:var(--surface);
  border:1px solid var(--border);
  border-radius:var(--r-lg);
  overflow:hidden;
  display:flex;flex-direction:column;
  min-height:600px;
}
.doc-toolbar{
  display:flex;align-items:center;gap:var(--s-2);
  padding:var(--s-2) var(--s-3);
  border-bottom:1px solid var(--border-soft);
  background:var(--surface);
}
.seg{
  display:inline-flex;
  background:var(--surface-2);
  border:1px solid var(--border-soft);
  border-radius:8px;
  padding:2px;gap:0;
}
.seg button{
  padding:5px 12px;
  font-size:12px;font-weight:500;
  color:var(--text-2);
  background:transparent;border:none;border-radius:6px;
  cursor:pointer;transition:all .15s;
}
.seg button:hover{color:var(--text)}
.seg button.active{
  background:var(--surface);
  color:var(--text);font-weight:600;
  box-shadow:var(--shadow-xs);
}
.toolbar-divider{width:1px;height:20px;background:var(--border-soft)}
.toolbar-spacer{flex:1}
.icon-btn{
  display:inline-flex;align-items:center;gap:6px;
  padding:5px 10px;
  font-size:12px;color:var(--text-2);font-weight:500;
  background:transparent;border:1px solid transparent;border-radius:6px;
  cursor:pointer;transition:all .15s;
}
.icon-btn:hover{background:var(--surface-2);color:var(--text)}
.icon-btn svg{width:14px;height:14px}
.icon-btn.toggle-on{background:var(--brand-soft);color:var(--brand);border-color:var(--brand-ring)}

.doc-scroll{
  flex:1;overflow-y:auto;
  background:var(--surface-warm);
  padding:var(--s-5) var(--s-4);
}
.doc-page{
  background:var(--surface);
  border:1px solid var(--border-soft);
  border-radius:var(--r-xl);
  padding:var(--s-5) var(--s-6);
  max-width:820px;margin:0 auto;
  font-size:17px;line-height:1.75;
  color:#1e293b;
  white-space:pre-wrap;word-wrap:break-word;
  box-shadow:var(--shadow-sm);
  font-feature-settings:"liga","calt","tnum";
}
.doc-page.focus-mode{color:rgba(30,41,59,.32)}
.doc-page.focus-mode .pii{color:#1e293b}

/* ----- PII spans (softer treatment) ----- */
.pii{
  position:relative;
  padding:1px 4px;
  border-radius:4px;
  cursor:pointer;
  transition:all .15s ease;
  background:var(--pii-tint,rgba(0,0,0,.04));
  color:var(--pii-text,inherit);
  box-shadow:inset 0 -1.5px 0 var(--pii-color,transparent);
}
.pii:hover{
  background:var(--pii-color,#888);
  color:#fff;
}
.pii.dimmed{
  opacity:.22;
  background:transparent;
  box-shadow:none;
  color:inherit;
}
.pii.linked{
  box-shadow:0 0 0 2px var(--pii-color,#888), inset 0 -1.5px 0 var(--pii-color,transparent);
  background:var(--pii-color,#888);color:#fff;
}

.mask-token{
  display:inline-block;
  padding:1px 8px;
  border-radius:4px;
  font-family:'JetBrains Mono',ui-monospace,monospace;
  font-size:13px;font-weight:600;
  background:var(--pii-tint,rgba(0,0,0,.06));
  color:var(--pii-text,inherit);
  border:1px dashed var(--pii-color,#888);
  letter-spacing:.02em;
}

/* ----- sidebar (inspection rail) ----- */
.rail{
  background:var(--surface);
  border:1px solid var(--border);
  border-radius:var(--r-lg);
  display:flex;flex-direction:column;
  overflow:hidden;
  max-height:calc(100vh - 260px);
}
.rail-section{
  padding:var(--s-3);
  border-bottom:1px solid var(--border-soft);
}
.rail-section:last-child{border-bottom:none}
.rail-section-header{
  display:flex;align-items:center;justify-content:space-between;
  margin-bottom:var(--s-2);
}
.rail-section-header .section-label{margin:0}
.rail-count{
  font-size:11px;color:var(--text-3);
  font-variant-numeric:tabular-nums;font-weight:600;
}

.findings{display:flex;flex-direction:column;gap:2px}
.finding{
  display:flex;align-items:center;gap:var(--s-2);
  padding:8px 10px;
  border-radius:8px;
  cursor:pointer;user-select:none;
  transition:all .15s;
  border:1px solid transparent;
}
.finding:hover{background:var(--surface-2)}
.finding.linked{background:var(--pii-tint);border-color:var(--pii-color)}
.finding.inactive{opacity:.42}
.finding-dot{width:8px;height:8px;border-radius:50%;flex-shrink:0;background:var(--pii-color,#888)}
.finding-label{flex:1;font-size:13.5px;font-weight:500;color:var(--text)}
.finding-count{
  font-size:12px;font-weight:600;color:var(--text-2);
  background:var(--surface-2);
  padding:2px 8px;border-radius:999px;
  font-variant-numeric:tabular-nums;
  border:1px solid var(--border-soft);
}
.finding-toggle{
  position:relative;
  width:26px;height:14px;border-radius:999px;
  background:var(--border);
  transition:background .18s;flex-shrink:0;
}
.finding.active .finding-toggle{background:var(--pii-color,var(--brand))}
.finding-toggle::after{
  content:'';position:absolute;
  top:1.5px;left:1.5px;
  width:11px;height:11px;border-radius:50%;
  background:#fff;
  transition:transform .18s;
  box-shadow:0 1px 2px rgba(0,0,0,.15);
}
.finding.active .finding-toggle::after{transform:translateX(12px)}

.rail-actions{display:flex;flex-direction:column;gap:6px}
.rail-actions .btn{justify-content:flex-start;width:100%}
.rail-actions .btn-ghost{font-size:13px;padding:8px 12px}

/* speakers */
.speakers{display:flex;flex-direction:column;gap:2px}
.speaker{
  display:flex;align-items:center;gap:var(--s-2);
  padding:6px 10px;border-radius:6px;cursor:pointer;
  font-size:13px;color:var(--text);
  transition:background .15s;
}
.speaker:hover{background:var(--surface-2)}
.speaker-name{flex:1;font-weight:500}
.speaker-count{
  font-size:11px;color:var(--text-3);font-weight:600;
  font-variant-numeric:tabular-nums;
}

/* empty state */
.empty-state{
  color:var(--text-3);font-size:13px;text-align:center;
  padding:var(--s-3);
}

/* tooltip */
.tooltip{
  position:fixed;
  background:rgba(15,23,42,.95);
  color:#fff;
  padding:6px 10px;border-radius:6px;
  font-size:12px;font-weight:500;
  pointer-events:none;z-index:999;
  white-space:nowrap;max-width:320px;overflow:hidden;text-overflow:ellipsis;
  box-shadow:var(--shadow-md);
  backdrop-filter:blur(8px);
}
.tooltip .tt-label{
  font-size:10px;text-transform:uppercase;letter-spacing:.06em;
  opacity:.7;margin-right:6px;
}

/* loading */
#loading{
  position:fixed;inset:0;
  background:rgba(247,247,249,.86);
  backdrop-filter:blur(10px);
  display:none;flex-direction:column;align-items:center;justify-content:center;
  z-index:9999;gap:var(--s-2);
}
.spinner{
  width:40px;height:40px;
  border:3px solid var(--border);
  border-top-color:var(--brand);
  border-radius:50%;
  animation:spin .7s linear infinite;
}
@keyframes spin{to{transform:rotate(360deg)}}
#loading .load-title{font-size:14px;font-weight:600;color:var(--text);margin-top:var(--s-2)}
#loading .load-sub{font-size:12px;color:var(--text-3)}

.error-banner{
  background:#fef2f2;border:1px solid #fecaca;
  color:#991b1b;
  padding:var(--s-2) var(--s-3);
  border-radius:var(--r-sm);
  margin:var(--s-3) var(--s-4);
  font-size:13px;display:none;
}

@media(max-width:960px){
  .metrics{grid-template-columns:repeat(2,1fr)}
  .layout{grid-template-columns:1fr}
  .rail{max-height:none}
  .features{grid-template-columns:1fr}
}
</style>
</head>
<body>

<!-- ============ UPLOAD VIEW ============ -->
<div id="upload-view">
  <div class="upload-card">
    <div class="brand">
      <div class="brand-logo">PR</div>
      <span class="brand-name">PII Reveal</span>
    </div>
    <div class="upload-hero">Document privacy inspection</div>
    <div class="upload-sub">Upload a document to detect and review sensitive content.</div>

    <div class="dropzone" id="dropzone">
      <div class="dropzone-icon">
        <svg width="22" height="22" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
          <path d="M12 3v12"/><path d="m7 8 5-5 5 5"/><path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"/>
        </svg>
      </div>
      <div class="dropzone-text">Drop your document, or click to browse</div>
      <div class="dropzone-hint">PDF, DOC, DOCX &middot; up to 128k tokens</div>
      <input type="file" id="file-input" accept=".pdf,.doc,.docx">
    </div>

    <div class="features">
      <div class="feature"><div class="feature-title">8 entity types</div><div class="feature-desc">Names, addresses, emails, phones, URLs, dates, accounts, secrets</div></div>
      <div class="feature"><div class="feature-title">128k context</div><div class="feature-desc">Full documents analyzed in a single pass</div></div>
      <div class="feature"><div class="feature-title">Context-aware</div><div class="feature-desc">Distinguishes &ldquo;May&rdquo; as a name vs. a month</div></div>
    </div>
    <div class="powered-by">Powered by <strong>OpenAI Privacy Filter</strong> &middot; Apache 2.0</div>
  </div>
</div>

<!-- ============ RESULTS VIEW ============ -->
<div id="results-view">

  <!-- header -->
  <header class="topbar">
    <div class="brand">
      <div class="brand-logo">PR</div>
      <span class="brand-name">PII Reveal</span>
    </div>
    <div class="file-pill" id="file-pill">
      <svg width="13" height="13" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
        <path d="M14 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V8z"/><polyline points="14 2 14 8 20 8"/>
      </svg>
      <span class="file-pill-name" id="file-name"></span>
    </div>
    <span class="topbar-summary" id="topbar-summary"></span>
    <div class="topbar-spacer"></div>
    <div class="topbar-actions">
      <button class="btn btn-ghost" id="btn-export-json">
        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"/><polyline points="7 10 12 15 17 10"/><line x1="12" x2="12" y1="15" y2="3"/></svg>
        Export JSON
      </button>
      <button class="btn btn-ghost" id="btn-copy-masked">
        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><rect width="14" height="14" x="8" y="8" rx="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>
        Copy masked
      </button>
      <button class="btn btn-primary" id="btn-new">New file</button>
    </div>
  </header>

  <div class="error-banner" id="error-banner"></div>

  <!-- KPI cards -->
  <section class="metrics">
    <div class="metric">
      <div class="metric-label">Sensitive content</div>
      <div class="metric-value" id="m-pct">0%</div>
      <div class="metric-hint" id="m-pct-hint">of document characters</div>
    </div>
    <div class="metric">
      <div class="metric-label">Detected entities</div>
      <div class="metric-value" id="m-spans">0</div>
      <div class="metric-hint" id="m-spans-hint">spans flagged</div>
    </div>
    <div class="metric">
      <div class="metric-label">Entity types</div>
      <div class="metric-value" id="m-cats">0</div>
      <div class="metric-hint" id="m-cats-hint">of 8 categories</div>
    </div>
    <div class="metric">
      <div class="metric-label">Risk level</div>
      <div class="metric-value metric-risk" id="m-risk">
        <span class="dot"></span><span class="lvl">&mdash;</span>
      </div>
      <div class="metric-hint" id="m-risk-hint">based on density &amp; type</div>
    </div>
  </section>

  <!-- legend + distribution -->
  <section class="legend">
    <div class="legend-header">
      <span class="section-label">Detected categories</span>
      <span class="section-label" id="legend-total" style="color:var(--text-3)"></span>
    </div>
    <div class="dist-bar" id="dist-bar"></div>
    <div class="chips" id="chips"></div>
  </section>

  <!-- main layout -->
  <div class="layout">
    <!-- document viewer -->
    <main class="doc-shell">
      <div class="doc-toolbar">
        <div class="seg" id="view-seg">
          <button data-view="original" class="active">Original</button>
          <button data-view="masked">Masked</button>
        </div>
        <div class="toolbar-divider"></div>
        <button class="icon-btn" id="btn-focus" title="Dim everything except entities">
          <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><circle cx="12" cy="12" r="3"/><path d="M2 12s3-7 10-7 10 7 10 7-3 7-10 7-10-7-10-7Z"/></svg>
          Focus mode
        </button>
        <div class="toolbar-spacer"></div>
        <button class="icon-btn" id="btn-prev" title="Previous entity">
          <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="m15 18-6-6 6-6"/></svg>
        </button>
        <span id="nav-pos" style="font-size:12px;color:var(--text-3);font-variant-numeric:tabular-nums;min-width:52px;text-align:center">0 / 0</span>
        <button class="icon-btn" id="btn-next" title="Next entity">
          <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="m9 18 6-6-6-6"/></svg>
        </button>
      </div>
      <div class="doc-scroll" id="doc-scroll">
        <div class="doc-page" id="doc-page"></div>
      </div>
    </main>

    <!-- inspection rail -->
    <aside class="rail">
      <div class="rail-section">
        <div class="rail-section-header">
          <span class="section-label">Findings</span>
          <span class="rail-count" id="findings-count"></span>
        </div>
        <div class="findings" id="findings"></div>
      </div>

      <div class="rail-section" id="speakers-section" style="display:none">
        <div class="rail-section-header">
          <span class="section-label">Speakers</span>
          <span class="rail-count" id="speakers-count"></span>
        </div>
        <div class="speakers" id="speakers"></div>
      </div>

      <div class="rail-section">
        <div class="rail-section-header">
          <span class="section-label">Actions</span>
        </div>
        <div class="rail-actions">
          <button class="btn btn-ghost" id="act-select-all">Select all categories</button>
          <button class="btn btn-ghost" id="act-clear-all">Clear selection</button>
          <button class="btn btn-ghost" id="act-copy-masked-2">Copy masked text</button>
          <button class="btn btn-ghost" id="act-export-json">Export findings (JSON)</button>
        </div>
      </div>
    </aside>
  </div>
</div>

<!-- loading -->
<div id="loading">
  <div class="spinner"></div>
  <div class="load-title">Analyzing document</div>
  <div class="load-sub">OpenAI Privacy Filter &middot; 128k context</div>
</div>

<div class="tooltip" id="tooltip" style="display:none"></div>

<script>
/* ===== state ===== */
const S = {
  text:'', spans:[], stats:{}, speakers:{}, catMeta:{}, filename:'',
  activeCats:new Set(), activeSpeakers:new Set(),
  view:'original',            // 'original' | 'masked'
  focus:false,
  sortedSpans:[], visibleIdx:[], navPos:-1,
};

/* ===== labels (fallback when backend meta missing) ===== */
const LBL = {private_person:'Person',private_address:'Address',private_email:'Email',private_phone:'Phone',private_url:'URL',private_date:'Date',account_number:'Account',secret:'Secret'};
const COL = {private_person:'#dc2626',private_address:'#0891b2',private_email:'#2563eb',private_phone:'#16a34a',private_url:'#ca8a04',private_date:'#9333ea',account_number:'#ea580c',secret:'#b91c1c'};
const TINT = {private_person:'rgba(220,38,38,.08)',private_address:'rgba(8,145,178,.08)',private_email:'rgba(37,99,235,.08)',private_phone:'rgba(22,163,74,.08)',private_url:'rgba(202,138,4,.10)',private_date:'rgba(147,51,234,.08)',account_number:'rgba(234,88,12,.08)',secret:'rgba(185,28,28,.10)'};
const TEXT = {private_person:'#991b1b',private_address:'#155e75',private_email:'#1e40af',private_phone:'#14532d',private_url:'#713f12',private_date:'#6b21a8',account_number:'#7c2d12',secret:'#7f1d1d'};
const ORDER = ['private_person','private_email','private_phone','private_address','private_date','private_url','account_number','secret'];

const metaFor = (c) => {
  const m = S.catMeta[c] || {};
  return {
    color: m.color || COL[c] || '#64748b',
    tint:  m.tint  || TINT[c] || 'rgba(100,116,139,.08)',
    text:  m.text  || TEXT[c] || '#334155',
    label: m.label || LBL[c] || c,
  };
};

/* ===== upload flow ===== */
const dz = document.getElementById('dropzone');
const fi = document.getElementById('file-input');
['dragenter','dragover'].forEach(e => dz.addEventListener(e, ev => { ev.preventDefault(); dz.classList.add('dragover'); }));
['dragleave','drop'].forEach(e => dz.addEventListener(e, ev => { ev.preventDefault(); dz.classList.remove('dragover'); }));
dz.addEventListener('drop', ev => { if (ev.dataTransfer.files[0]) uploadFile(ev.dataTransfer.files[0]); });
fi.addEventListener('change', ev => { if (ev.target.files[0]) uploadFile(ev.target.files[0]); });

async function uploadFile(file) {
  const ext = file.name.split('.').pop().toLowerCase();
  if (!['pdf','doc','docx'].includes(ext)) { showError('Unsupported file type.'); return; }
  document.getElementById('loading').style.display = 'flex';
  document.getElementById('upload-view').style.display = 'none';
  const form = new FormData(); form.append('file', file);
  try {
    const r = await fetch('/api/analyze', { method:'POST', body: form });
    const d = await r.json();
    if (d.error) { showError(d.error); return; }
    S.text = d.text; S.spans = d.spans; S.stats = d.stats;
    S.speakers = d.speakers || {}; S.catMeta = d.categories_meta || {};
    S.filename = d.filename;
    S.activeCats = new Set(Object.keys(d.stats.categories));
    S.activeSpeakers = new Set(Object.keys(S.speakers));
    S.sortedSpans = [...S.spans].sort((a,b) => a.start - b.start);
    S.navPos = -1;
    renderResults();
  } catch (e) {
    showError('Analysis failed: ' + e.message);
  } finally {
    document.getElementById('loading').style.display = 'none';
  }
}

function showError(m) {
  document.getElementById('loading').style.display = 'none';
  document.getElementById('results-view').style.display = 'block';
  const b = document.getElementById('error-banner');
  b.textContent = m; b.style.display = 'block';
}

function resetView() {
  document.getElementById('results-view').style.display = 'none';
  document.getElementById('upload-view').style.display = 'flex';
  document.getElementById('error-banner').style.display = 'none';
  fi.value = '';
}
document.getElementById('btn-new').addEventListener('click', resetView);

/* ===== render ===== */
function renderResults() {
  document.getElementById('results-view').style.display = 'block';
  document.getElementById('error-banner').style.display = 'none';
  document.getElementById('file-name').textContent = S.filename;
  document.getElementById('topbar-summary').textContent =
    `${S.stats.total_spans} entities across ${S.stats.num_categories} categories`;
  renderMetrics();
  renderLegend();
  renderFindings();
  renderSpeakers();
  renderDoc();
  updateNavPos();
}

function renderMetrics() {
  const s = S.stats;
  document.getElementById('m-pct').textContent = (s.pii_percentage ?? 0) + '%';
  document.getElementById('m-pct-hint').textContent = `${s.pii_chars.toLocaleString()} of ${s.total_chars.toLocaleString()} chars`;
  document.getElementById('m-spans').textContent = s.total_spans;
  document.getElementById('m-spans-hint').textContent = 'spans flagged';
  document.getElementById('m-cats').textContent = s.num_categories;
  document.getElementById('m-cats-hint').textContent = 'of 8 possible';
  const risk = (s.risk_level || 'Low').toLowerCase();
  const rEl = document.getElementById('m-risk');
  rEl.className = 'metric-value metric-risk ' + risk;
  rEl.querySelector('.lvl').textContent = s.risk_level || 'Low';
}

function renderLegend() {
  const s = S.stats, total = s.total_chars || 1;
  // distribution bar — only fills the fraction covered by PII
  const bar = document.getElementById('dist-bar');
  bar.innerHTML = '';
  const ordered = ORDER.filter(c => s.categories[c]);
  for (const c of ordered) {
    const info = s.categories[c], m = metaFor(c);
    const seg = document.createElement('div');
    seg.className = 'dist-seg';
    seg.style.width = (info.chars / total * 100) + '%';
    seg.style.background = m.color;
    seg.dataset.cat = c;
    seg.title = `${m.label} · ${info.count} spans · ${info.chars} chars`;
    seg.addEventListener('mouseenter', () => highlightCategory(c, true));
    seg.addEventListener('mouseleave', () => highlightCategory(c, false));
    bar.appendChild(seg);
  }
  document.getElementById('legend-total').textContent = `${s.total_spans} entities`;
  // chips
  const ch = document.getElementById('chips'); ch.innerHTML = '';
  for (const c of ordered) {
    const info = s.categories[c], m = metaFor(c);
    const el = document.createElement('span');
    el.className = 'chip';
    const active = S.activeCats.has(c);
    if (!active) el.classList.add('inactive');
    el.style.background = m.tint;
    el.style.color = m.text;
    el.style.borderColor = 'transparent';
    el.innerHTML = `<span class="chip-dot" style="background:${m.color}"></span>${m.label}<span class="chip-count">${info.count}</span>`;
    el.addEventListener('click', () => toggleCategory(c));
    el.addEventListener('mouseenter', () => highlightCategory(c, true));
    el.addEventListener('mouseleave', () => highlightCategory(c, false));
    ch.appendChild(el);
  }
}

function renderFindings() {
  const box = document.getElementById('findings');
  box.innerHTML = '';
  const cats = S.stats.categories;
  const ordered = ORDER.filter(c => cats[c]);
  document.getElementById('findings-count').textContent = `${ordered.length} types`;
  if (!ordered.length) { box.innerHTML = '<div class="empty-state">No entities detected.</div>'; return; }
  for (const c of ordered) {
    const m = metaFor(c), info = cats[c], active = S.activeCats.has(c);
    const el = document.createElement('div');
    el.className = 'finding' + (active ? ' active' : ' inactive');
    el.dataset.cat = c;
    el.style.setProperty('--pii-color', m.color);
    el.style.setProperty('--pii-tint', m.tint);
    el.innerHTML = `
      <span class="finding-dot"></span>
      <span class="finding-label">${m.label}</span>
      <span class="finding-count">${info.count}</span>
      <span class="finding-toggle"></span>`;
    el.addEventListener('click', () => toggleCategory(c));
    el.addEventListener('mouseenter', () => highlightCategory(c, true));
    el.addEventListener('mouseleave', () => highlightCategory(c, false));
    box.appendChild(el);
  }
}

function renderSpeakers() {
  const sec = document.getElementById('speakers-section'), box = document.getElementById('speakers');
  const names = Object.keys(S.speakers);
  if (!names.length) { sec.style.display = 'none'; return; }
  sec.style.display = 'block';
  document.getElementById('speakers-count').textContent = `${names.length}`;
  box.innerHTML = '';
  for (const name of names) {
    const el = document.createElement('div');
    el.className = 'speaker';
    el.innerHTML = `<span class="speaker-name">${esc(name)}</span><span class="speaker-count">${S.speakers[name]}</span>`;
    box.appendChild(el);
  }
}

function esc(s) { const d = document.createElement('div'); d.textContent = s; return d.innerHTML; }

function renderDoc() {
  const { text, sortedSpans, view, activeCats, focus } = S;
  const page = document.getElementById('doc-page');
  page.classList.toggle('focus-mode', focus);

  let html = '', pos = 0;
  S.visibleIdx = [];
  for (let i = 0; i < sortedSpans.length; i++) {
    const sp = sortedSpans[i];
    if (sp.start < pos) continue;
    if (sp.start > pos) html += esc(text.substring(pos, sp.start));
    const m = metaFor(sp.label);
    const on = activeCats.has(sp.label);
    if (on) S.visibleIdx.push(i);
    const style = `--pii-color:${m.color};--pii-tint:${m.tint};--pii-text:${m.text}`;
    if (view === 'masked' && on) {
      html += `<span class="mask-token" style="${style}" data-idx="${i}" data-cat="${sp.label}">[${m.label.toUpperCase()}]</span>`;
    } else {
      html += `<span class="pii${on ? '' : ' dimmed'}" style="${style}" data-idx="${i}" data-cat="${sp.label}" data-text="${esc(sp.text)}">${esc(text.substring(sp.start, sp.end))}</span>`;
    }
    pos = sp.end;
  }
  if (pos < text.length) html += esc(text.substring(pos));
  page.innerHTML = html;

  const tt = document.getElementById('tooltip');
  page.querySelectorAll('.pii, .mask-token').forEach(el => {
    const cat = el.dataset.cat, m = metaFor(cat);
    el.addEventListener('mouseenter', ev => {
      const orig = S.sortedSpans[+el.dataset.idx];
      tt.innerHTML = `<span class="tt-label">${m.label}</span>${esc(orig.text)}`;
      tt.style.display = 'block'; moveTT(ev);
      highlightCategory(cat, true);
    });
    el.addEventListener('mousemove', moveTT);
    el.addEventListener('mouseleave', () => {
      tt.style.display = 'none';
      highlightCategory(cat, false);
    });
  });

  if (S.navPos >= S.visibleIdx.length) S.navPos = -1;
}

function moveTT(ev) {
  const t = document.getElementById('tooltip');
  const w = t.offsetWidth || 200;
  const left = Math.min(ev.clientX + 12, window.innerWidth - w - 12);
  t.style.left = left + 'px';
  t.style.top = (ev.clientY - 34) + 'px';
}

/* ===== interactions ===== */
function toggleCategory(c) {
  if (S.activeCats.has(c)) S.activeCats.delete(c);
  else S.activeCats.add(c);
  renderLegend(); renderFindings(); renderDoc(); updateNavPos();
}

function highlightCategory(c, on) {
  // span side
  document.querySelectorAll('.pii, .mask-token').forEach(el => {
    if (el.dataset.cat === c) el.classList.toggle('linked', on);
  });
  // sidebar side
  document.querySelectorAll('.finding').forEach(el => {
    if (el.dataset.cat === c) el.classList.toggle('linked', on);
  });
}

/* selection bulk actions */
document.getElementById('act-select-all').addEventListener('click', () => {
  S.activeCats = new Set(Object.keys(S.stats.categories));
  renderLegend(); renderFindings(); renderDoc(); updateNavPos();
});
document.getElementById('act-clear-all').addEventListener('click', () => {
  S.activeCats = new Set();
  renderLegend(); renderFindings(); renderDoc(); updateNavPos();
});

/* view segmented control */
document.getElementById('view-seg').addEventListener('click', ev => {
  const btn = ev.target.closest('button[data-view]');
  if (!btn) return;
  S.view = btn.dataset.view;
  document.querySelectorAll('#view-seg button').forEach(b => b.classList.toggle('active', b === btn));
  renderDoc();
});

/* focus mode */
document.getElementById('btn-focus').addEventListener('click', ev => {
  S.focus = !S.focus;
  ev.currentTarget.classList.toggle('toggle-on', S.focus);
  renderDoc();
});

/* next/prev */
document.getElementById('btn-next').addEventListener('click', () => navigate(1));
document.getElementById('btn-prev').addEventListener('click', () => navigate(-1));
function navigate(dir) {
  if (!S.visibleIdx.length) return;
  S.navPos = (S.navPos + dir + S.visibleIdx.length) % S.visibleIdx.length;
  const i = S.visibleIdx[S.navPos];
  const el = document.querySelector(`[data-idx="${i}"]`);
  if (el) {
    el.scrollIntoView({ behavior:'smooth', block:'center' });
    el.classList.add('linked');
    setTimeout(() => el.classList.remove('linked'), 1200);
  }
  updateNavPos();
}
function updateNavPos() {
  const total = S.visibleIdx.length;
  const cur = S.navPos >= 0 ? (S.navPos + 1) : 0;
  document.getElementById('nav-pos').textContent = `${cur} / ${total}`;
}

/* export + copy */
function maskedText() {
  const parts = []; let pos = 0;
  for (const sp of S.sortedSpans) {
    if (sp.start < pos) continue;
    if (sp.start > pos) parts.push(S.text.substring(pos, sp.start));
    const m = metaFor(sp.label);
    parts.push(S.activeCats.has(sp.label) ? `[${m.label.toUpperCase()}]` : S.text.substring(sp.start, sp.end));
    pos = sp.end;
  }
  if (pos < S.text.length) parts.push(S.text.substring(pos));
  return parts.join('');
}

function download(name, content, type = 'application/json') {
  const blob = new Blob([content], { type });
  const a = document.createElement('a');
  a.href = URL.createObjectURL(blob); a.download = name;
  document.body.appendChild(a); a.click(); a.remove();
  setTimeout(() => URL.revokeObjectURL(a.href), 1000);
}

function exportJSON() {
  download((S.filename || 'findings') + '.findings.json',
    JSON.stringify({ filename:S.filename, stats:S.stats, spans:S.spans, speakers:S.speakers }, null, 2));
}

async function copyMasked() {
  try {
    await navigator.clipboard.writeText(maskedText());
    flashButton('Copied');
  } catch { flashButton('Copy failed'); }
}

let _flashTimer;
function flashButton(msg) {
  const b = document.getElementById('btn-copy-masked');
  const prev = b.innerHTML;
  b.innerHTML = msg;
  clearTimeout(_flashTimer);
  _flashTimer = setTimeout(() => { b.innerHTML = prev; }, 1200);
}

document.getElementById('btn-export-json').addEventListener('click', exportJSON);
document.getElementById('act-export-json').addEventListener('click', exportJSON);
document.getElementById('btn-copy-masked').addEventListener('click', copyMasked);
document.getElementById('act-copy-masked-2').addEventListener('click', copyMasked);

/* keyboard: n / p for next/prev, f for focus */
document.addEventListener('keydown', ev => {
  if (document.getElementById('results-view').style.display === 'none') return;
  if (ev.target.matches('input,textarea')) return;
  if (ev.key === 'n' || ev.key === 'ArrowDown') { ev.preventDefault(); navigate(1); }
  else if (ev.key === 'p' || ev.key === 'ArrowUp') { ev.preventDefault(); navigate(-1); }
  else if (ev.key === 'f') { document.getElementById('btn-focus').click(); }
});
</script>
</body>
</html>"""

# ── launch ───────────────────────────────────────────────────────
if __name__ == "__main__":
    server.launch(server_name="0.0.0.0", server_port=7860)