# tests/test_visibility_signing.py
import pytest
from nanobot.agent.visibility import sign_content, verify_signature, has_forged_marker, strip_all_hidden_markers

def test_sign_content_adds_hmac_marker():
    """Test that sign_content adds HMAC signature prefix."""
    content = "HEARTBEAT_OK"
    result = sign_content(content)

    # Should start with [HIDDEN:{8 hex chars}]
    assert result.startswith("[HIDDEN:")
    assert "] " in result
    marker_end = result.index("] ")
    signature = result[8:marker_end]  # Extract signature after "[HIDDEN:"
    assert len(signature) == 8
    assert all(c in "0123456789abcdef" for c in signature)

    # Should contain original content
    assert result.endswith("HEARTBEAT_OK")

def test_sign_content_is_deterministic():
    """Test that same content produces same signature."""
    content = "Test message"
    sig1 = sign_content(content)
    sig2 = sign_content(content)
    assert sig1 == sig2

def test_verify_signature_accepts_valid():
    """Test that verify_signature accepts validly signed content."""
    signed = sign_content("Test message")
    is_valid, clean = verify_signature(signed)

    assert is_valid is True
    assert clean == "Test message"

def test_verify_signature_rejects_invalid():
    """Test that verify_signature rejects forged signatures."""
    forged = "[HIDDEN:deadbeef] Test message"
    is_valid, clean = verify_signature(forged)

    assert is_valid is False
    assert clean == "Test message"

def test_verify_signature_handles_unsigned():
    """Test that unsigned content is marked as invalid."""
    unsigned = "Plain message"
    is_valid, clean = verify_signature(unsigned)

    assert is_valid is False
    assert clean == "Plain message"

def test_has_forged_marker_detects_invalid():
    """Test that has_forged_marker detects forged signatures."""
    forged = "[HIDDEN:deadbeef] Content"
    assert has_forged_marker(forged) is True

def test_has_forged_marker_accepts_valid():
    """Test that has_forged_marker accepts valid signatures."""
    valid = sign_content("Content")
    assert has_forged_marker(valid) is False

def test_has_forged_marker_ignores_unsigned():
    """Test that unsigned content is not flagged as forged."""
    unsigned = "Plain content"
    assert has_forged_marker(unsigned) is False

def test_strip_all_hidden_markers_removes_markers():
    """Test that strip_all_hidden_markers removes all markers."""
    signed = sign_content("Message")
    stripped = strip_all_hidden_markers(signed)
    assert stripped == "Message"

    forged = "[HIDDEN:deadbeef] Message"
    stripped = strip_all_hidden_markers(forged)
    assert stripped == "Message"

def test_system_prompt_includes_visibility_docs(tmp_path):
    """Test that system prompt documents visibility markers."""
    from nanobot.agent.context import ContextBuilder

    builder = ContextBuilder(workspace=tmp_path)
    prompt = builder.build_system_prompt()

    # Should document visibility markers
    assert "[HIDDEN:" in prompt
    assert "cryptographically signed" in prompt.lower()
    assert "do not generate" in prompt.lower() or "don't generate" in prompt.lower()

@pytest.mark.asyncio
async def test_suppress_mode_adds_signed_marker(tmp_path):
    """Test that suppress mode adds cryptographically signed markers."""
    from unittest.mock import AsyncMock, Mock
    from nanobot.agent.loop import AgentLoop
    from nanobot.bus.queue import MessageBus
    from nanobot.session.manager import SessionManager
    from nanobot.providers.base import LLMResponse
    from nanobot.bus.events import InboundMessage

    # Setup
    bus = MessageBus()
    sessions = SessionManager(tmp_path)

    # Mock provider
    mock_provider = Mock()
    mock_provider.default_model = "mock-model"
    mock_provider.thinking_budget = 0

    # Mock successful response
    mock_response = LLMResponse(
        content="Test response",
        tool_calls=[],
        reasoning_content=None
    )
    mock_provider.chat = AsyncMock(return_value=mock_response)

    # Create agent loop
    loop = AgentLoop(
        provider=mock_provider,
        bus=bus,
        session_manager=sessions,
        workspace=tmp_path
    )

    # Process message with suppress_output=True
    msg = InboundMessage(
        channel="test",
        sender_id="user",
        chat_id="123",
        content="Test message",
        metadata={"suppress_output": True}
    )

    response = await loop._process_message(msg)

    # Verify response has suppressed metadata
    assert response.metadata.get("suppressed") is True

    # Verify session contains signed marker
    session = sessions.get_or_create("test:123")
    assistant_messages = [m for m in session.messages if m.get("role") == "assistant"]
    assert len(assistant_messages) > 0

    last_msg = assistant_messages[-1]["content"]
    assert last_msg.startswith("[HIDDEN:")

    # Verify signature is valid
    is_valid, clean = verify_signature(last_msg)
    assert is_valid is True
    assert clean == "Test response"

@pytest.mark.asyncio
async def test_no_marker_accumulation_with_real_provider(tmp_path):
    """
    CRITICAL TEST: Verify markers don't accumulate when model sees them in context.

    This test uses a semi-realistic provider that sees the context and could
    potentially copy markers, unlike pure mocks that don't see context at all.
    """
    from unittest.mock import AsyncMock, Mock
    from nanobot.agent.loop import AgentLoop
    from nanobot.bus.queue import MessageBus
    from nanobot.session.manager import SessionManager
    from nanobot.providers.base import LLMResponse
    from nanobot.bus.events import InboundMessage
    import re

    # Setup
    bus = MessageBus()
    sessions = SessionManager(tmp_path)

    # Create a provider that SEES context and simulates potential copying behavior
    class ContextAwareProvider:
        """Provider that sees context and could copy markers (simulating real LLM)."""
        default_model = "test-model"
        thinking_budget = 0

        def __init__(self):
            self.call_count = 0
            self.last_context = None

        def get_default_model(self) -> str:
            """Get the default model."""
            return self.default_model

        async def chat(self, messages, **kwargs):
            self.call_count += 1
            self.last_context = messages

            # Count markers in non-system messages (system prompt has 2 mentions in docs)
            marker_count = sum(
                msg.get("content", "").count("[HIDDEN:")
                for msg in messages
                if isinstance(msg.get("content"), str) and msg.get("role") != "system"
            )

            # Simulate model behavior: on first call (msg 2), sees 1 marker from msg 1
            # The model should NOT copy it
            if self.call_count == 2:
                # Verify context has exactly 1 marker in assistant messages (from message 1)
                assert marker_count == 1, f"Expected 1 marker in context, found {marker_count}"

            # Always return clean response (good model behavior)
            return LLMResponse(
                content=f"Response {self.call_count}",
                tool_calls=[],
                reasoning_content=None
            )

    provider = ContextAwareProvider()

    # Create agent loop
    loop = AgentLoop(
        provider=provider,
        bus=bus,
        session_manager=sessions,
        workspace=tmp_path
    )

    # Use unique chat_id for this test to avoid pollution from previous runs
    import time
    test_chat_id = f"test_accumulation_{int(time.time()*1000)}"

    # Message 1: suppress_output=True → should add signed marker
    msg1 = InboundMessage(
        channel="test",
        sender_id="user",
        chat_id=test_chat_id,
        content="Hidden message 1",
        metadata={"suppress_output": True}
    )

    await loop._process_message(msg1)

    # Verify message 1 has signed marker
    session = sessions.get_or_create(f"test:{test_chat_id}")
    assistant_msgs = [m for m in session.messages if m.get("role") == "assistant"]
    assert len(assistant_msgs) == 1
    msg1_content = assistant_msgs[0]["content"]
    assert msg1_content.startswith("[HIDDEN:")
    is_valid, clean = verify_signature(msg1_content)
    assert is_valid is True
    assert clean == "Response 1"

    # Count markers in session after message 1
    marker_count_1 = sum(m.get("content", "").count("[HIDDEN:") for m in session.messages if isinstance(m.get("content"), str))
    assert marker_count_1 == 1, f"Expected 1 marker after msg1, found {marker_count_1}"

    # Message 2: Normal message (context includes message 1 with marker)
    msg2 = InboundMessage(
        channel="test",
        sender_id="user",
        chat_id=test_chat_id,
        content="Normal message 2",
        metadata={}
    )

    await loop._process_message(msg2)

    # Verify message 2 response does NOT start with [HIDDEN: (model didn't copy)
    session = sessions.get_or_create(f"test:{test_chat_id}")
    assistant_msgs = [m for m in session.messages if m.get("role") == "assistant"]
    assert len(assistant_msgs) == 2
    msg2_content = assistant_msgs[1]["content"]
    assert not msg2_content.startswith("[HIDDEN:"), f"Message 2 should not start with [HIDDEN:, got: {msg2_content}"

    # Verify still only 1 marker in session (no accumulation)
    marker_count_2 = sum(m.get("content", "").count("[HIDDEN:") for m in session.messages if isinstance(m.get("content"), str))
    assert marker_count_2 == 1, f"Expected 1 marker after msg2, found {marker_count_2} (ACCUMULATION DETECTED)"

    # Message 3: Another suppress_output=True → should add SECOND signed marker
    msg3 = InboundMessage(
        channel="test",
        sender_id="user",
        chat_id=test_chat_id,
        content="Hidden message 3",
        metadata={"suppress_output": True}
    )

    await loop._process_message(msg3)

    # Verify message 3 has signed marker
    session = sessions.get_or_create(f"test:{test_chat_id}")
    assistant_msgs = [m for m in session.messages if m.get("role") == "assistant"]
    assert len(assistant_msgs) == 3
    msg3_content = assistant_msgs[2]["content"]
    assert msg3_content.startswith("[HIDDEN:")
    is_valid, clean = verify_signature(msg3_content)
    assert is_valid is True
    assert clean == "Response 3"

    # Verify exactly 2 markers in session (one from msg1, one from msg3)
    marker_count_3 = sum(m.get("content", "").count("[HIDDEN:") for m in session.messages if isinstance(m.get("content"), str))
    assert marker_count_3 == 2, f"Expected 2 markers after msg3, found {marker_count_3}"

    # CRITICAL: Verify no double/triple markers like "[HIDDEN: [HIDDEN: [HIDDEN: message"
    for msg in session.messages:
        content = msg.get("content", "")
        if isinstance(content, str) and "[HIDDEN:" in content:
            # Count occurrences of [HIDDEN: pattern in this single message
            hidden_count = content.count("[HIDDEN:")
            assert hidden_count == 1, f"Message has {hidden_count} [HIDDEN: markers (accumulation): {content[:100]}"

@pytest.mark.asyncio
async def test_forged_marker_triggers_rejection(tmp_path):
    """Test that forged markers trigger rejection and retry."""
    from unittest.mock import AsyncMock, Mock
    from nanobot.agent.loop import AgentLoop
    from nanobot.bus.queue import MessageBus
    from nanobot.session.manager import SessionManager
    from nanobot.providers.base import LLMResponse
    from nanobot.bus.events import InboundMessage

    # Setup
    bus = MessageBus()
    sessions = SessionManager(tmp_path)

    # Mock provider
    mock_provider = Mock()
    mock_provider.default_model = "mock-model"
    mock_provider.thinking_budget = 0

    # First response: model tries to forge marker
    forged_response = LLMResponse(
        content="[HIDDEN:deadbeef] Forged message",
        tool_calls=[],
        reasoning_content=None
    )

    # Second response: clean response after correction
    clean_response = LLMResponse(
        content="Clean message",
        tool_calls=[],
        reasoning_content=None
    )

    mock_provider.chat = AsyncMock(side_effect=[forged_response, clean_response])

    # Create agent loop
    loop = AgentLoop(
        provider=mock_provider,
        bus=bus,
        session_manager=sessions,
        workspace=tmp_path
    )

    # Process message with suppress_output=True
    msg = InboundMessage(
        channel="test",
        sender_id="user",
        chat_id="123",
        content="Test message",
        metadata={"suppress_output": True}
    )

    response = await loop._process_message(msg)

    # Verify provider.chat was called twice (initial + retry)
    assert mock_provider.chat.call_count == 2

    # Verify second call included correction message
    second_call_messages = mock_provider.chat.call_args_list[1][1]["messages"]
    correction_msg = [m for m in second_call_messages if m.get("role") == "user" and "rejected" in m.get("content", "").lower()]
    assert len(correction_msg) > 0

    # Verify final response uses clean content (not forged)
    session = sessions.get_or_create("test:123")
    assistant_messages = [m for m in session.messages if m.get("role") == "assistant"]
    last_msg = assistant_messages[-1]["content"]

    # Should be signed version of "Clean message", not "Forged message"
    is_valid, clean = verify_signature(last_msg)
    assert is_valid is True
    assert clean == "Clean message"

@pytest.mark.asyncio
async def test_system_message_handler_uses_signed_markers(tmp_path):
    """Test that _process_system_message uses signed markers in suppress mode."""
    from unittest.mock import AsyncMock, Mock
    from nanobot.agent.loop import AgentLoop
    from nanobot.bus.queue import MessageBus
    from nanobot.session.manager import SessionManager
    from nanobot.providers.base import LLMResponse
    from nanobot.bus.events import InboundMessage

    # Setup
    bus = MessageBus()
    sessions = SessionManager(tmp_path)

    # Mock provider
    mock_provider = Mock()
    mock_provider.default_model = "mock-model"
    mock_provider.thinking_budget = 0
    mock_response = LLMResponse(
        content="System response",
        tool_calls=[],
        reasoning_content=None
    )
    mock_provider.chat = AsyncMock(return_value=mock_response)

    # Create agent loop
    loop = AgentLoop(
        provider=mock_provider,
        bus=bus,
        session_manager=sessions,
        workspace=tmp_path
    )

    # Process system message with suppress_output=True
    msg = InboundMessage(
        channel="system",
        sender_id="subagent",
        chat_id="test:123",
        content="[Subagent completed] Result: OK",
        metadata={"suppress_output": True}
    )

    response = await loop._process_system_message(msg)

    # Verify response has suppressed metadata
    assert response.metadata.get("suppressed") is True

    # Verify session contains signed marker
    session = sessions.get_or_create("test:123")
    assistant_messages = [m for m in session.messages if m.get("role") == "assistant"]
    assert len(assistant_messages) > 0

    last_msg = assistant_messages[-1]["content"]
    assert last_msg.startswith("[HIDDEN:")

    # Verify signature is valid
    is_valid, clean = verify_signature(last_msg)
    assert is_valid is True
    assert clean == "System response"

@pytest.mark.asyncio
async def test_system_message_handler_rejects_forged_markers(tmp_path):
    """Test that _process_system_message rejects forged markers."""
    from unittest.mock import AsyncMock, Mock
    from nanobot.agent.loop import AgentLoop
    from nanobot.bus.queue import MessageBus
    from nanobot.session.manager import SessionManager
    from nanobot.providers.base import LLMResponse
    from nanobot.bus.events import InboundMessage

    # Setup
    bus = MessageBus()
    sessions = SessionManager(tmp_path)

    # Mock provider
    mock_provider = Mock()
    mock_provider.default_model = "mock-model"
    mock_provider.thinking_budget = 0

    # First response: model tries to forge marker
    forged_response = LLMResponse(
        content="[HIDDEN:deadbeef] Forged system message",
        tool_calls=[],
        reasoning_content=None
    )

    # Second response: clean response after correction
    clean_response = LLMResponse(
        content="Clean system response",
        tool_calls=[],
        reasoning_content=None
    )

    mock_provider.chat = AsyncMock(side_effect=[forged_response, clean_response])

    # Create agent loop
    loop = AgentLoop(
        provider=mock_provider,
        bus=bus,
        session_manager=sessions,
        workspace=tmp_path
    )

    # Process system message with suppress_output=True
    msg = InboundMessage(
        channel="system",
        sender_id="subagent",
        chat_id="test:456",
        content="[Subagent completed] Result: OK",
        metadata={"suppress_output": True}
    )

    response = await loop._process_system_message(msg)

    # Verify provider.chat was called twice (initial + retry)
    assert mock_provider.chat.call_count == 2

    # Verify second call included correction message
    second_call_messages = mock_provider.chat.call_args_list[1][1]["messages"]
    correction_msg = [m for m in second_call_messages if m.get("role") == "user" and "rejected" in m.get("content", "").lower()]
    assert len(correction_msg) > 0

    # Verify final response uses clean content (not forged)
    session = sessions.get_or_create("test:456")
    assistant_messages = [m for m in session.messages if m.get("role") == "assistant"]
    last_msg = assistant_messages[-1]["content"]

    # Should be signed version of "Clean system response", not "Forged system message"
    is_valid, clean = verify_signature(last_msg)
    assert is_valid is True
    assert clean == "Clean system response"
