"""Session-scoped pseudonymization.

Assigns stable identifiers (PERSON_1, EMAIL_1, …) per unique PII value
within a session, enabling coherent multi-turn LLM conversations.
The session can later restore pseudonyms back to originals in the response.
"""
from dataclasses import dataclass, field
from typing import Dict, List, Union


# Type labels that get short human-readable prefixes in pseudonyms.
_TYPE_PREFIX: Dict[str, str] = {
    "NAME": "PERSON",
    "EMAIL": "EMAIL",
    "PHONE": "PHONE",
    "DATE_OF_BIRTH": "DOB",
    "CREDIT_CARD": "CARD",
    "IBAN": "ACCOUNT",
    "PESEL": "ID",
    "BR_CPF": "ID",
    "PT_NIF": "ID",
    "PT_NISS": "ID",
    "SECRET": "TOKEN",
    "UUID": "ID",
    "ADDRESS": "ADDRESS",
    "URL": "URL",
    "LONG_NUMERIC_ID": "ID",
    "ORG": "ORG",
}


@dataclass
class PseudonymSession:
    """Maintains a bidirectional mapping of original PII ↔ pseudonym.

    Thread-safety: not thread-safe; create one session per request/conversation.
    """
    _forward: Dict[str, str] = field(default_factory=dict)   # original → pseudonym
    _reverse: Dict[str, str] = field(default_factory=dict)   # pseudonym → original
    _counters: Dict[str, int] = field(default_factory=dict)  # prefix → count

    def pseudonym_for(self, original: str, pii_type: str) -> str:
        """Return the stable pseudonym for *original*, creating one if needed."""
        if original in self._forward:
            return self._forward[original]

        prefix = _TYPE_PREFIX.get(pii_type, "PII")
        count = self._counters.get(prefix, 0) + 1
        self._counters[prefix] = count
        pseudonym = f"{prefix}_{count}"

        self._forward[original] = pseudonym
        self._reverse[pseudonym] = original
        return pseudonym

    def restore(self, text: str) -> str:
        """Replace all pseudonyms in *text* with their original values."""
        if not self._reverse:
            return text
        # Sort by length descending so longer tokens are replaced first
        # (avoids partial matches on similar-prefix pseudonyms like TOKEN_1/TOKEN_10)
        for pseudonym in sorted(self._reverse, key=len, reverse=True):
            text = text.replace(pseudonym, self._reverse[pseudonym])
        return text

    @property
    def mappings(self) -> Dict[str, str]:
        """Return a copy of the forward mapping (original → pseudonym)."""
        return dict(self._forward)

    def clear(self) -> None:
        """Reset the session."""
        self._forward.clear()
        self._reverse.clear()
        self._counters.clear()


def apply_pseudonym_redaction(
    text: str,
    findings: List[Dict[str, Union[int, str]]],
    session: PseudonymSession,
) -> str:
    """Replace detected PII spans with session-stable pseudonyms.

    Conflict resolution (same span, different types) is handled upstream
    by _resolve_conflicts; here we trust the findings are already deduplicated.
    """
    from upf.redaction.text import _resolve_conflicts

    ordered = _resolve_conflicts(findings)

    out: List[str] = []
    cursor = 0
    for item in ordered:
        start = int(item["start"])
        end = int(item["end"])

        if start < cursor or end <= start:
            continue

        out.append(text[cursor:start])
        original = str(item.get("text", text[start:end]))
        out.append(session.pseudonym_for(original, str(item["type"])))
        cursor = end

    out.append(text[cursor:])
    return "".join(out)
