"""
Tests for llm-taint.

Key scenarios:
- TaintedStr behaves as str
- taint() labels values correctly
- check_sink() raises on unsafe sinks
- scrub() / scrub_dict() redacts tainted values
- taint_env_secrets() populates _env_taint_registry
- get_tainted_env() returns TaintedStr (survives os.environ POSIX round-trip)
- TaintScrubFilter scrubs log records
- Runtime sink/key registration (add_safe_sink, add_unsafe_sink, add_secret_env_key)
"""
from __future__ import annotations

import logging
import os

import pytest

from llm_taint import (
    TaintViolationError,
    TaintedStr,
    _env_taint_registry,
    add_safe_sink,
    add_secret_env_key,
    add_unsafe_sink,
    check_sink,
    get_label,
    get_tainted_env,
    is_tainted,
    scrub,
    scrub_dict,
    taint,
    taint_env_secrets,
)


# ============================================================================
# TaintedStr behaviour
# ============================================================================

class TestTaintedStr:
    def test_is_str(self):
        t = taint("sk-abc123", label="test_key")
        assert isinstance(t, str)
        assert t == "sk-abc123"

    def test_label_stored(self):
        t = taint("my-secret", label="my_label")
        assert t._taint_label == "my_label"
        assert get_label(t) == "my_label"

    def test_plain_str_not_tainted(self):
        assert not is_tainted("plain")
        assert get_label("plain") is None

    def test_is_tainted_true(self):
        t = taint("value", label="lbl")
        assert is_tainted(t)

    def test_concatenation_preserves_taint(self):
        t = taint("sk-", label="key")
        result = t + "suffix"
        assert isinstance(result, TaintedStr)
        assert result == "sk-suffix"
        assert result._taint_label == "key"

    def test_slice_loses_taint(self):
        """Slices intentionally drop taint to avoid accidental propagation."""
        t = taint("sk-abc", label="key")
        sliced = t[0:3]
        assert not isinstance(sliced, TaintedStr)
        assert sliced == "sk-"

    def test_repr_does_not_expose_value(self):
        t = taint("super-secret-value", label="key")
        r = repr(t)
        assert "super-secret-value" not in r
        assert "key" in r

    def test_equality_with_plain_str(self):
        t = taint("hello", label="lbl")
        assert t == "hello"
        assert "hello" == t

    def test_length(self):
        t = taint("abc", label="lbl")
        assert len(t) == 3

    def test_none_not_tainted(self):
        assert not is_tainted(None)
        assert get_label(None) is None


# ============================================================================
# check_sink
# ============================================================================

class TestCheckSink:
    def test_raises_on_log(self):
        t = taint("sk-secret", label="api_key")
        with pytest.raises(TaintViolationError) as exc_info:
            check_sink(t, sink="log")
        assert exc_info.value.label == "api_key"
        assert exc_info.value.sink == "log"

    def test_raises_on_http_response(self):
        t = taint("tok-secret", label="token")
        with pytest.raises(TaintViolationError):
            check_sink(t, sink="http_response")

    def test_raises_on_tool_result(self):
        t = taint("sk-secret", label="key")
        with pytest.raises(TaintViolationError):
            check_sink(t, sink="tool_result")

    def test_raises_on_websocket(self):
        t = taint("sk-secret", label="key")
        with pytest.raises(TaintViolationError):
            check_sink(t, sink="websocket")

    def test_allows_vault(self):
        t = taint("sk-secret", label="api_key")
        check_sink(t, sink="vault")  # no raise

    def test_allows_encrypted(self):
        t = taint("sk-secret", label="api_key")
        check_sink(t, sink="encrypted")  # no raise

    def test_allows_llm_prompt(self):
        t = taint("sk-secret", label="api_key")
        check_sink(t, sink="llm_prompt")  # no raise

    def test_allows_unknown_sink(self):
        t = taint("sk-secret", label="api_key")
        check_sink(t, sink="some_custom_sink")  # fail-open — no raise

    def test_plain_str_does_not_raise(self):
        check_sink("just a normal string", sink="log")  # no raise

    def test_violation_error_message(self):
        t = taint("sk-secret", label="mykey")
        with pytest.raises(TaintViolationError) as exc_info:
            check_sink(t, sink="log")
        msg = str(exc_info.value)
        assert "mykey" in msg
        assert "log" in msg


# ============================================================================
# scrub / scrub_dict
# ============================================================================

class TestScrub:
    def test_scrub_tainted(self):
        t = taint("sk-supersecret", label="api_key")
        assert scrub(t) == "[REDACTED:api_key]"

    def test_scrub_plain(self):
        assert scrub("hello world") == "hello world"

    def test_scrub_int(self):
        assert scrub(42) == "42"

    def test_scrub_dict_tainted_values(self):
        t = taint("sk-secret", label="key")
        result = scrub_dict({"safe": "value", "secret": t})
        assert result["safe"] == "value"
        assert result["secret"] == "[REDACTED:key]"

    def test_scrub_dict_plain_values(self):
        result = scrub_dict({"a": "hello", "b": "world"})
        assert result == {"a": "hello", "b": "world"}


# ============================================================================
# taint_env_secrets + get_tainted_env
# ============================================================================

class TestEnvTainting:
    def test_taint_env_secrets_taints_known_keys(self):
        env = {
            "OPENAI_API_KEY": "sk-real-key-value-1234",
            "SOME_PLAIN_VAR": "not-a-secret",
        }
        result = taint_env_secrets(env)
        assert isinstance(result["OPENAI_API_KEY"], TaintedStr)
        assert not isinstance(result["SOME_PLAIN_VAR"], TaintedStr)

    def test_taint_env_secrets_anthropic(self):
        env = {"ANTHROPIC_API_KEY": "sk-ant-test-key-1234567890"}
        result = taint_env_secrets(env)
        assert isinstance(result["ANTHROPIC_API_KEY"], TaintedStr)

    def test_taint_env_secrets_populates_registry(self):
        env = {"OPENAI_API_KEY": "sk-test-key-for-registry-1234"}
        taint_env_secrets(env)
        assert "OPENAI_API_KEY" in _env_taint_registry
        assert isinstance(_env_taint_registry["OPENAI_API_KEY"], TaintedStr)

    def test_taint_env_secrets_empty_value_not_tainted(self):
        env = {"OPENAI_API_KEY": ""}
        result = taint_env_secrets(env)
        assert not isinstance(result["OPENAI_API_KEY"], TaintedStr)

    def test_get_tainted_env_returns_tainted_str(self, monkeypatch):
        """get_tainted_env() returns TaintedStr even if os.environ returns plain str."""
        monkeypatch.setenv("OPENAI_API_KEY", "sk-test-value-posix-fix-1234")
        taint_env_secrets({"OPENAI_API_KEY": "sk-test-value-posix-fix-1234"})

        result = get_tainted_env("OPENAI_API_KEY")
        assert isinstance(result, TaintedStr), (
            "get_tainted_env must return TaintedStr — os.environ strips subclass on POSIX"
        )
        assert result == "sk-test-value-posix-fix-1234"

    def test_get_tainted_env_falls_back_to_os_environ(self, monkeypatch):
        monkeypatch.setenv("SOME_NON_SECRET_VAR", "hello")
        _env_taint_registry.pop("SOME_NON_SECRET_VAR", None)
        result = get_tainted_env("SOME_NON_SECRET_VAR")
        assert result == "hello"
        assert not isinstance(result, TaintedStr)

    def test_get_tainted_env_default(self, monkeypatch):
        monkeypatch.delenv("NO_SUCH_KEY_XYZ", raising=False)
        _env_taint_registry.pop("NO_SUCH_KEY_XYZ", None)
        assert get_tainted_env("NO_SUCH_KEY_XYZ", default="fallback") == "fallback"

    def test_posix_strips_tainted_str(self, monkeypatch):
        """Demonstrates the POSIX limitation that the registry fix works around."""
        t = taint("sk-test-value", label="demo")
        monkeypatch.setitem(os.environ, "DEMO_TAINT_KEY_TEST", t)
        retrieved = os.environ.get("DEMO_TAINT_KEY_TEST")
        assert retrieved == "sk-test-value"
        assert not isinstance(retrieved, TaintedStr), (
            "Confirms: os.environ strips TaintedStr on POSIX"
        )

    def test_custom_secret_key(self):
        add_secret_env_key("MY_CUSTOM_SECRET")
        env = {"MY_CUSTOM_SECRET": "super-secret-value-abc123"}
        result = taint_env_secrets(env)
        assert isinstance(result["MY_CUSTOM_SECRET"], TaintedStr)


# ============================================================================
# Runtime sink registration
# ============================================================================

class TestSinkRegistration:
    def test_add_safe_sink(self):
        t = taint("sk-secret", label="key")
        add_safe_sink("my_safe_sink")
        check_sink(t, sink="my_safe_sink")  # no raise

    def test_add_unsafe_sink(self):
        t = taint("sk-secret", label="key")
        add_unsafe_sink("my_unsafe_sink")
        with pytest.raises(TaintViolationError):
            check_sink(t, sink="my_unsafe_sink")


# ============================================================================
# TaintScrubFilter
# ============================================================================

class TestTaintScrubFilter:
    def test_scrubs_tainted_args_tuple(self):
        from llm_taint.logger import TaintScrubFilter

        filt = TaintScrubFilter()
        record = logging.LogRecord(
            name="test", level=logging.INFO,
            pathname="", lineno=0,
            msg="key=%s",
            args=(taint("sk-secret", label="key"),),
            exc_info=None,
        )
        filt.filter(record)
        assert record.args[0] == "[REDACTED:key]"

    def test_scrubs_tainted_args_dict(self):
        from llm_taint.logger import TaintScrubFilter

        filt = TaintScrubFilter()
        # Construct with None args, then assign dict directly to avoid
        # Python 3.12 LogRecord.__init__ subscripting the dict with int key 0
        record = logging.LogRecord(
            name="test", level=logging.INFO,
            pathname="", lineno=0,
            msg="key=%(key)s",
            args=None,
            exc_info=None,
        )
        record.args = {"key": taint("sk-secret", label="mykey")}
        filt.filter(record)
        assert record.args["key"] == "[REDACTED:mykey]"

    def test_passes_plain_args(self):
        from llm_taint.logger import TaintScrubFilter

        filt = TaintScrubFilter()
        record = logging.LogRecord(
            name="test", level=logging.INFO,
            pathname="", lineno=0,
            msg="value=%s",
            args=("plain_value",),
            exc_info=None,
        )
        filt.filter(record)
        assert record.args[0] == "plain_value"

    def test_scrubs_tainted_msg(self):
        from llm_taint.logger import TaintScrubFilter

        filt = TaintScrubFilter()
        record = logging.LogRecord(
            name="test", level=logging.INFO,
            pathname="", lineno=0,
            msg=taint("sk-secret", label="msg_key"),
            args=(),
            exc_info=None,
        )
        filt.filter(record)
        assert record.msg == "[REDACTED:msg_key]"

    def test_filter_always_returns_true(self):
        from llm_taint.logger import TaintScrubFilter

        filt = TaintScrubFilter()
        record = logging.LogRecord(
            name="test", level=logging.INFO,
            pathname="", lineno=0, msg="ok", args=(), exc_info=None,
        )
        assert filt.filter(record) is True

    def test_install_is_idempotent(self):
        from llm_taint.logger import install_taint_filter
        install_taint_filter()
        install_taint_filter()  # second call — no error, no double install
