"""
Taint tracking — lightweight secret flow analysis for LLM pipelines.

Labels sensitive values (API keys, tokens, passwords) at their entry point
and raises TaintViolationError if a tainted value reaches an unsafe sink
(logs, HTTP responses, error messages, tool outputs).

Usage::

    from llm_taint import taint, check_sink, TaintedStr

    # Label a secret at entry
    api_key = taint(os.environ["OPENAI_API_KEY"], label="openai_api_key")

    # Check before logging (raises TaintViolationError if tainted)
    check_sink(some_value, sink="log")

    # TaintedStr behaves exactly like str but carries a label
    assert isinstance(api_key, str)  # True — transparent to most code

The POSIX problem::

    On Linux/macOS, os.environ encodes values to bytes and decodes back to
    plain str on every read, stripping the TaintedStr subclass.  Use
    taint_env_secrets() + get_tainted_env() to work around this.
"""
from __future__ import annotations

import logging
import re
from typing import Optional, Set

logger = logging.getLogger(__name__)

# Patterns that look like secrets even without explicit taint labeling
_SECRET_PATTERNS: tuple[re.Pattern, ...] = (
    re.compile(r"sk-[A-Za-z0-9]{20,}"),            # OpenAI / Anthropic keys
    re.compile(r"sk-ant-[A-Za-z0-9\-]{20,}"),      # Anthropic keys
    re.compile(r"Bearer [A-Za-z0-9\-._~+/]+=*"),    # Bearer tokens
    re.compile(r"[A-Za-z0-9]{32,}"),                # Generic long tokens (heuristic)
    re.compile(r"password\s*[:=]\s*\S+", re.I),    # password= patterns
    re.compile(r"api.?key\s*[:=]\s*\S+", re.I),    # api_key= patterns
    re.compile(r"secret\s*[:=]\s*\S+", re.I),      # secret= patterns
)

# Safe sinks (allowed to receive tainted data)
_SAFE_SINKS: frozenset[str] = frozenset({
    "llm_prompt",   # LLM receives secrets only when explicitly authorized
    "vault",        # Vault storage is the correct place for secrets
    "encrypted",    # Encrypted transport is safe
})

# Unsafe sinks where tainted values must not appear
_UNSAFE_SINKS: frozenset[str] = frozenset({
    "log",
    "http_response",
    "tool_result",
    "error_message",
    "websocket",
})


class TaintViolationError(RuntimeError):
    """Raised when a tainted value reaches an unsafe sink."""

    def __init__(self, label: str, sink: str, preview: str = ""):
        self.label = label
        self.sink = sink
        self.preview = preview
        super().__init__(
            f"TAINT VIOLATION: secret '{label}' reached unsafe sink '{sink}'"
            + (f" — value preview: {preview[:20]}..." if preview else "")
        )


class TaintedStr(str):
    """
    A ``str`` subclass that carries a taint label.

    Behaves identically to ``str`` in all string operations.
    The taint label is preserved through concatenation when possible
    (Python ``str`` subclass limitations apply — slices intentionally
    lose taint to avoid unexpected propagation).

    Do not construct directly — use :func:`taint` instead.
    """

    _taint_label: str

    def __new__(cls, value: str, label: str) -> "TaintedStr":
        instance = super().__new__(cls, value)
        instance._taint_label = label
        return instance

    def __repr__(self) -> str:
        return f"TaintedStr(label={self._taint_label!r}, len={len(self)})"

    def __add__(self, other: str) -> "TaintedStr":  # type: ignore[override]
        result = super().__add__(other)
        return TaintedStr(result, self._taint_label)

    def __getitem__(self, key) -> str:  # type: ignore[override]
        # Slices lose taint intentionally
        return str.__getitem__(self, key)


# Internal registry — keyed by id() of each TaintedStr
_tainted_ids: Set[int] = set()

# Registry of tainted env vars — populated by taint_env_secrets().
# Exists because os.environ on POSIX encodes values to bytes and decodes
# back to plain str on every read, stripping the TaintedStr subclass.
# Use get_tainted_env() instead of os.environ[key] for sensitive values.
_env_taint_registry: dict[str, "TaintedStr"] = {}


def taint(value: str, label: str) -> TaintedStr:
    """
    Label a string value as tainted (sensitive).

    Returns a :class:`TaintedStr` that behaves like ``str`` but carries the
    label.  Call this at the point where a secret enters the system — at
    config load, environment variable read, or database retrieval.

    Args:
        value: The secret string to taint.
        label: A human-readable label (e.g. ``"openai_api_key"``, ``"db_password"``).

    Returns:
        :class:`TaintedStr` instance.

    Example::

        import os
        from llm_taint import taint

        api_key = taint(os.environ["OPENAI_API_KEY"], label="openai_api_key")
        assert isinstance(api_key, str)   # True — transparent to normal code
        assert api_key == os.environ["OPENAI_API_KEY"]
    """
    t = TaintedStr(value, label)
    _tainted_ids.add(id(t))
    logger.debug("TAINT: labeled '%s' (%d chars)", label, len(value))
    return t


def is_tainted(value: object) -> bool:
    """Return ``True`` if the value is a :class:`TaintedStr`."""
    return isinstance(value, TaintedStr)


def get_label(value: object) -> Optional[str]:
    """Return the taint label if the value is tainted, else ``None``."""
    if isinstance(value, TaintedStr):
        return value._taint_label
    return None


def check_sink(value: object, sink: str) -> None:
    """
    Check whether a value is safe to send to the given sink.

    Raises :class:`TaintViolationError` if the value is a :class:`TaintedStr`
    and the sink is in the unsafe sink set.  Also emits a warning (but does
    not raise) if a plain string heuristically looks like a secret.

    Built-in unsafe sinks: ``"log"``, ``"http_response"``, ``"tool_result"``,
    ``"error_message"``, ``"websocket"``.

    Built-in safe sinks: ``"llm_prompt"``, ``"vault"``, ``"encrypted"``.

    Unknown sinks are silently allowed (fail-open).

    Args:
        value: The value being sent to the sink.
        sink:  The sink name.

    Raises:
        TaintViolationError: If a tainted value reaches an unsafe sink.

    Example::

        from llm_taint import taint, check_sink

        key = taint(os.environ["OPENAI_API_KEY"], label="openai_api_key")

        # Before logging
        user_message = request.json["message"]
        check_sink(user_message, sink="log")  # safe if untainted

        # This would raise:
        check_sink(key, sink="log")  # TaintViolationError!
    """
    if sink in _SAFE_SINKS:
        return

    if sink not in _UNSAFE_SINKS:
        return  # Unknown sink — allow by default

    if isinstance(value, TaintedStr):
        raise TaintViolationError(value._taint_label, sink, str(value))

    # Heuristic check for unmarked secrets (defense in depth — warning only)
    if isinstance(value, str) and _looks_like_secret(value):
        logger.warning(
            "TAINT HEURISTIC: potential secret at sink '%s' (not explicitly tainted)",
            sink,
        )


def scrub(value: object) -> str:
    """
    Return a safe representation of a potentially tainted value.

    - Tainted → ``"[REDACTED:<label>]"``
    - Heuristically secret → ``"[REDACTED:heuristic]"``
    - Otherwise → ``str(value)``

    Use this before logging or serializing values that might be tainted.

    Example::

        from llm_taint import taint, scrub, scrub_dict

        key = taint("sk-abc123", label="openai_key")
        print(scrub(key))      # "[REDACTED:openai_key]"
        print(scrub("hello"))  # "hello"

        config = {"api_key": key, "model": "gpt-4"}
        safe = scrub_dict(config)
        print(safe)  # {"api_key": "[REDACTED:openai_key]", "model": "gpt-4"}
    """
    if isinstance(value, TaintedStr):
        return f"[REDACTED:{value._taint_label}]"
    s = str(value)
    if _looks_like_secret(s):
        return "[REDACTED:heuristic]"
    return s


def scrub_dict(d: dict) -> dict:
    """
    Return a copy of a dict with all tainted/secret-looking values scrubbed.

    Safe for logging or serialization.
    """
    return {k: scrub(v) for k, v in d.items()}


def _looks_like_secret(value: str) -> bool:
    """Heuristic: does this string look like a secret?"""
    if len(value) < 16:
        return False
    for pattern in _SECRET_PATTERNS[:3]:  # Use only the most reliable patterns
        if pattern.search(value):
            return True
    return False


# ---------------------------------------------------------------------------
# Common environment variable secret keys
# ---------------------------------------------------------------------------

_SECRET_ENV_KEYS: frozenset[str] = frozenset({
    # LLM / AI providers
    "OPENAI_API_KEY",
    "ANTHROPIC_API_KEY",
    "COHERE_API_KEY",
    "MISTRAL_API_KEY",
    "TOGETHER_API_KEY",
    "REPLICATE_API_TOKEN",
    # Cloud / infra
    "AWS_SECRET_ACCESS_KEY",
    "GCP_SERVICE_ACCOUNT_KEY",
    "AZURE_CLIENT_SECRET",
    "DIGITALOCEAN_TOKEN",
    # Databases
    "DATABASE_URL",
    "REDIS_URL",
    "MONGODB_URI",
    # Auth / web
    "SECRET_KEY",
    "JWT_SECRET",
    "WEBHOOK_SECRET",
    "SESSION_SECRET",
    # Payment
    "STRIPE_SECRET_KEY",
    "STRIPE_WEBHOOK_SECRET",
    # Communications
    "SENDGRID_API_KEY",
    "TWILIO_AUTH_TOKEN",
    "TWITTER_BEARER_TOKEN",
    "DISCORD_BOT_TOKEN",
    "TELEGRAM_BOT_TOKEN",
    "SLACK_BOT_TOKEN",
    # Source control / CI
    "GITHUB_TOKEN",
    "GITLAB_TOKEN",
})


def taint_env_secrets(env: dict) -> dict:
    """
    Return a copy of the env dict with known secret keys tainted.

    Also populates ``_env_taint_registry`` so :func:`get_tainted_env` returns
    :class:`TaintedStr` even after ``os.environ`` strips the subclass on POSIX.

    Call this on ``os.environ`` or a config dict at application startup::

        import os
        from llm_taint import taint_env_secrets, get_tainted_env

        taint_env_secrets(dict(os.environ))   # call once at startup

        # Later, use get_tainted_env instead of os.environ for sensitive vars
        key = get_tainted_env("OPENAI_API_KEY")
        assert isinstance(key, TaintedStr)    # True — even on Linux/macOS

    Args:
        env: Dictionary of environment variables or config values.

    Returns:
        New dict with secret values replaced by :class:`TaintedStr` instances.
    """
    result = {}
    for k, v in env.items():
        if k in _SECRET_ENV_KEYS and isinstance(v, str) and v:
            t = taint(v, label=k.lower())
            result[k] = t
            _env_taint_registry[k] = t
        else:
            result[k] = v
    return result


def get_tainted_env(key: str, default: Optional[str] = None) -> Optional[str]:
    """
    Get an environment variable, returning a :class:`TaintedStr` when tainted.

    Prefer this over ``os.environ[key]`` / ``os.getenv(key)`` for values that
    should carry taint labels.  On POSIX systems, ``os.environ`` strips the
    :class:`TaintedStr` subclass on every read — this function works around
    that by checking the in-process ``_env_taint_registry`` first.

    Falls back to ``os.environ`` if the key is not in the registry.

    Args:
        key:     Environment variable name.
        default: Returned if the key is absent from both registry and env.
    """
    import os as _os
    if key in _env_taint_registry:
        return _env_taint_registry[key]
    return _os.environ.get(key, default)


def add_safe_sink(sink: str) -> None:
    """Register an additional safe sink name at runtime."""
    global _SAFE_SINKS
    _SAFE_SINKS = _SAFE_SINKS | {sink}


def add_unsafe_sink(sink: str) -> None:
    """Register an additional unsafe sink name at runtime."""
    global _UNSAFE_SINKS
    _UNSAFE_SINKS = _UNSAFE_SINKS | {sink}


def add_secret_env_key(key: str) -> None:
    """Register an additional environment variable name as a secret."""
    global _SECRET_ENV_KEYS
    _SECRET_ENV_KEYS = _SECRET_ENV_KEYS | {key}
