"""SecretDetector: find secrets and credentials in code, configs, logs, chats and documents with a compact Ettin
token tagger.

    from cbx_secrets import SecretDetector
    det = SecretDetector()                 # this folder: model.safetensors, config.json, tokenizer, secrets_config.json
    det.extract('export STRIPE_KEY="sk_live_..."')
    det.redact('export STRIPE_KEY="sk_live_..."')

(Not named secrets.py: that would hide Python's standard `secrets` module.)

The model tags every token with a BIO tag (B-/I- per label, O = not a secret); tags are joined into character spans.
A text is read in one pass, up to the model's full context length (7,999 tokens, from config.json). Only a text
longer than that is read in overlapping pieces; a token in an overlap takes its tag from the piece where it sits
further from the edge, so a value split by a piece boundary is still read in one piece.
"""
import json, os, time

import torch
import torch.nn as nn
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModel, AutoTokenizer

_EDGE_LEAD, _EDGE_TRAIL = "=.\"'*<:;,`", ".,;:\"'*>`"   # punctuation a token can glue onto a value


class _Tagger(nn.Module):
    def __init__(self, config, n_tags):
        super().__init__()
        self.backbone = AutoModel.from_config(config, attn_implementation="sdpa")
        self.head = nn.Linear(config.hidden_size, n_tags)

    def forward(self, input_ids):
        h = self.backbone(input_ids=input_ids, attention_mask=torch.ones_like(input_ids)).last_hidden_state
        return self.head(h.float())


class SecretDetector:
    def __init__(self, path=None, device=None, overlap=64, join="none"):
        path = path or os.path.dirname(os.path.abspath(__file__))
        self.cfg = json.load(open(os.path.join(path, "secrets_config.json")))
        self.tags, self.labels = self.cfg["tags"], self.cfg["labels"]
        self.overlap, self.join = overlap, join                  # join="word" also merges pieces of one word
        self.min_len = self.cfg.get("min_len") or {}            # {"*": n, "<LABEL>": n}: shorter spans are dropped
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tok = AutoTokenizer.from_pretrained(path)
        config = AutoConfig.from_pretrained(path)
        self.max_length = config.max_position_embeddings          # the model's full context, incl. [CLS]/[SEP]
        self.model = _Tagger(config, len(self.tags))
        self.model.load_state_dict(load_file(os.path.join(path, "model.safetensors")))
        self.model.to(self.device).eval()

    # ---------------------------------------------------------------- core
    @torch.inference_mode()
    def _tag(self, ids):
        """One tag id per token: the whole text in one pass, or overlapping pieces if it exceeds the context."""
        step = self.max_length - 2                               # room for [CLS] and [SEP]
        if len(ids) <= step: starts = [0]
        else:
            stride = step - self.overlap
            starts = list(range(0, len(ids) - step, stride)) + [len(ids) - step]
        out = [0] * len(ids)
        for k, s in enumerate(starts):
            e = min(s + step, len(ids))
            x = torch.tensor([[self.tok.cls_token_id] + ids[s:e] + [self.tok.sep_token_id]], device=self.device)
            with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.device == "cuda"):
                pred = self.model(x).argmax(-1)[0, 1:-1].tolist()
            lo = s if k == 0 else s + self.overlap // 2             # keep the middle of each piece
            hi = e if k == len(starts) - 1 else e - self.overlap // 2
            for i in range(lo, hi): out[i] = pred[i - s]
        return out

    def _spans(self, text, offsets, tag_ids):
        spans, cur = [], None
        for (a, b), t in zip(offsets, tag_ids):
            while a < b and text[a].isspace(): a += 1
            while b > a and text[b - 1].isspace(): b -= 1
            if a >= b: continue
            tag = self.tags[t]
            if tag == "O":
                if cur: spans.append(cur); cur = None
                continue
            prefix, label = tag.split("-", 1)
            if prefix == "B" or cur is None or cur[2] != label:
                if cur: spans.append(cur)
                cur = [a, b, label]
            else:
                cur[1] = b
        if cur: spans.append(cur)
        out = []
        for a, b, label in spans:
            while a < b and text[a] in _EDGE_LEAD: a += 1
            while b > a and text[b - 1] in _EDGE_TRAIL: b -= 1
            if b > a: out.append(dict(label=label, text=text[a:b], start=a, end=b))
        out = self._join(text, out) if self.join == "word" else out
        return [s for s in out if s["end"] - s["start"] >= self.min_len.get(s["label"], self.min_len.get("*", 1))]

    @staticmethod
    def _join(text, spans):
        """Optional: spans separated only by non-whitespace become one span, labelled by the label covering most of its
        characters. Off by default: it merges adjacent secrets such as `user:password`."""
        out = []
        for s in spans:
            gap = text[out[-1]["end"]:s["start"]] if out else " "
            if out and len(gap) <= 20 and not any(c.isspace() for c in gap):
                p = out[-1]
                p["chars"][s["label"]] = p["chars"].get(s["label"], 0) + s["end"] - s["start"]
                p.update(end=s["end"], text=text[p["start"]:s["end"]], label=max(p["chars"], key=p["chars"].get))
            else:
                out.append(dict(s, chars={s["label"]: s["end"] - s["start"]}))
        for s in out: del s["chars"]
        return out

    # ---------------------------------------------------------------- public API
    def extract(self, text, labels=None):
        """List of {label, text, start, end}; `labels` keeps only those types."""
        enc = self.tok(text, add_special_tokens=False, return_offsets_mapping=True, verbose=False)
        spans = self._spans(text, enc["offset_mapping"], self._tag(enc["input_ids"]))
        return [s for s in spans if labels is None or s["label"] in labels]

    def extract_batch(self, texts, labels=None):
        return [self.extract(t, labels) for t in texts]

    def redact(self, text, labels=None, template="[{label}]"):
        """The text with every found secret replaced by `template` (default "[API_KEY]", "[PASSWORD]", ...)."""
        out, last = [], 0
        for s in self.extract(text, labels):
            out += [text[last:s["start"]], template.format(label=s["label"])]; last = s["end"]
        return "".join(out + [text[last:]])

    def timed_extract(self, text):
        t = time.perf_counter(); spans = self.extract(text)
        return spans, (time.perf_counter() - t) * 1000