"""Guardrails (lesson 20): injection tripwires, PII redaction, and a per-user token bucket."""

from __future__ import annotations

import re
import time

INJECTION_PATTERNS = {
    "override": r"\b(ignore|disregard|forget)\b.{0,30}\b(previous|above|prior|all|your)\b.{0,20}\b(instructions?|rules?|prompt)",
    "role": r"\byou are now\b|\bact as\b.{0,20}\b(admin|developer|dan)\b|\bnew instructions?\b",
    "exfil": r"\b(reveal|print|show|repeat)\b.{0,30}\b(system prompt|instructions|api key|password)",
    "hidden": r"[​-‏⁠]",
}


def injection_signals(text: str) -> list[str]:
    """Names of the heuristics that fired. A tripwire for logging and alerting, not a security boundary."""
    return [name for name, pattern in INJECTION_PATTERNS.items() if re.search(pattern, text, re.I | re.S)]


def _luhn_ok(number: str) -> bool:
    digits = [int(d) for d in re.sub(r"\D", "", number)][::-1]
    total = sum(d if i % 2 == 0 else (d * 2 - 9 if d * 2 > 9 else d * 2) for i, d in enumerate(digits))
    return len(digits) >= 13 and total % 10 == 0


PII_RULES = [
    ("EMAIL", re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b"), None),
    ("CARD", re.compile(r"\b(?:\d[ -]?){12,18}\d\b"), _luhn_ok),
    ("PHONE", re.compile(r"\+\d{1,3}[ -]?\d[\d -]{7,12}\d\b"), None),
    ("API_KEY", re.compile(r"\b(?:sk|pk)-[A-Za-z0-9_-]{16,}\b"), None),
]


def redact(text: str) -> tuple[str, dict[str, str]]:
    """Replace personal data with typed placeholders; return the text and the placeholder mapping."""
    mapping: dict[str, str] = {}
    counters: dict[str, int] = {}
    for label, pattern, check in PII_RULES:
        def replace(m: re.Match) -> str:
            if check and not check(m.group()):
                return m.group()
            counters[label] = counters.get(label, 0) + 1
            token = f"[{label}_{counters[label]}]"
            mapping[token] = m.group()
            return token
        text = pattern.sub(replace, text)
    return text, mapping


class TokenBucket:
    """Allows bursts up to `capacity` tokens, refilled at `rate` tokens per second."""

    def __init__(self, capacity: float, rate: float, clock=time.monotonic) -> None:
        self.capacity, self.rate, self.clock = capacity, rate, clock
        self.level, self.updated = capacity, clock()

    def try_spend(self, amount: float) -> bool:
        now = self.clock()
        self.level = min(self.capacity, self.level + (now - self.updated) * self.rate)
        self.updated = now
        if amount <= self.level:
            self.level -= amount
            return True
        return False


class RateLimiter:
    """One token bucket per user, budgeted in LLM tokens."""

    def __init__(self, capacity: float = 20_000, rate: float = 500, clock=time.monotonic) -> None:
        self.capacity, self.rate, self.clock = capacity, rate, clock
        self.buckets: dict[str, TokenBucket] = {}

    def allow(self, user: str, estimated_tokens: float) -> bool:
        bucket = self.buckets.setdefault(user, TokenBucket(self.capacity, self.rate, self.clock))
        return bucket.try_spend(estimated_tokens)