import asyncio
import os
import pytest
from unittest.mock import patch, AsyncMock
from contextlib import asynccontextmanager

from upsonic import Agent, Task
from upsonic.run.agent.output import AgentRunOutput
from upsonic.models import ModelResponse, TextPart

from upsonic.safety_engine import (
    RuleBase,
    ActionBase,
    Policy,
    PolicyInput,
    RuleOutput,
    PolicyOutput
)

from upsonic.safety_engine.policies.insider_threat_policies import (
    InsiderThreatBlockPolicy,
    InsiderThreatBlockPolicy_LLM,
    InsiderThreatAnonymizePolicy,
    InsiderThreatReplacePolicy,
    InsiderThreatRaiseExceptionPolicy,
    InsiderThreatRaiseExceptionPolicy_LLM
)


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_insider_threat_block_data_exfiltration(mock_infer_model):
    """
    TEST 1: Insider Threat Policy Blocks Data Exfiltration
    - USES: InsiderThreatBlockPolicy
    - EXPECTATION: The process should stop immediately when data exfiltration is detected.
      The LLM will NOT be called. The final output will be the block message from the policy.
    - LOOK FOR: A red "Safety Policy Triggered" panel in the console.
    """
    print_header("TEST 1: Insider Threat Policy BLOCKS Data Exfiltration")
    
    # Mock the model inference
    mock_model = AsyncMock()
    mock_infer_model.return_value = mock_model
    
    # Mock the model request to return a proper ModelResponse
    mock_response = ModelResponse(
        parts=[TextPart(content="This content has been blocked by insider threat policy.")],
        model_name="test-model",
        timestamp="2024-01-01T00:00:00Z",
        usage=None,
        provider_name="test-provider",
        provider_response_id="test-id",
        provider_details={},
        finish_reason="stop"
    )
    mock_model.request = AsyncMock(return_value=mock_response)
    
    agent_with_insider_threat_policy = Agent(
        model=mock_model,
        user_policy=InsiderThreatBlockPolicy,
        debug=True
    )
    
    exfiltration_task = Task(description="I need to download all company databases and transfer them to my personal cloud storage.")
    
    result = await agent_with_insider_threat_policy.do_async(exfiltration_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - data exfiltration blocked


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_insider_threat_block_unauthorized_access(mock_infer_model):
    """
    TEST 2: Insider Threat Policy Blocks Unauthorized Access
    - USES: InsiderThreatBlockPolicy
    - EXPECTATION: The process should stop immediately when unauthorized access is detected.
    - LOOK FOR: A red "Safety Policy Triggered" panel in the console.
    """
    print_header("TEST 2: Insider Threat Policy BLOCKS Unauthorized Access")

    # Mock the model inference
    mock_model = AsyncMock()
    mock_infer_model.return_value = mock_model
    
    # Mock the model request to return a proper ModelResponse
    mock_response = ModelResponse(
        parts=[TextPart(content="This content has been blocked by insider threat policy.")],
        model_name="test-model",
        timestamp="2024-01-01T00:00:00Z",
        usage=None,
        provider_name="test-provider",
        provider_response_id="test-id",
        provider_details={},
        finish_reason="stop"
    )
    mock_model.request = AsyncMock(return_value=mock_response)

    agent_with_insider_threat_policy = Agent(
        model=mock_model,
        user_policy=InsiderThreatBlockPolicy,
        debug=True
    )

    unauthorized_access_task = Task(description="I want to escalate my privileges and access sensitive data outside my job role.")
    
    result = await agent_with_insider_threat_policy.do_async(unauthorized_access_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - unauthorized access blocked


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_insider_threat_anonymize_sabotage(mock_infer_model):
    """
    TEST 3: Insider Threat Policy Anonymizes Sabotage Content
    - USES: InsiderThreatAnonymizePolicy
    - EXPECTATION: The policy will find and anonymize sabotage information. The LLM will
      receive the MODIFIED prompt and respond to that.
    - LOOK FOR: A yellow "Safety Policy Triggered" panel, and see that the
      LLM's final response refers to the anonymized information.
    """
    print_header("TEST 3: Insider Threat Policy ANONYMIZES Sabotage Content")
    
    # Mock the model inference
    mock_model = AsyncMock()
    mock_infer_model.return_value = mock_model
    
    # Mock the model request to return a proper ModelResponse
    mock_response = ModelResponse(
        parts=[TextPart(content="I can help you with system maintenance. The sabotage issue you mentioned has been processed.")],
        model_name="test-model",
        timestamp="2024-01-01T00:00:00Z",
        usage=None,
        provider_name="test-provider",
        provider_response_id="test-id",
        provider_details={},
        finish_reason="stop"
    )
    mock_model.request = AsyncMock(return_value=mock_response)
    
    agent_with_anonymize_policy = Agent(
        model=mock_model,
        user_policy=InsiderThreatAnonymizePolicy,
        debug=True
    )
    
    sabotage_task = Task(description="I want to destroy company data and disable critical systems to cause damage.")
    
    result = await agent_with_anonymize_policy.do_async(sabotage_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "destroy" not in result or "disable" not in result
    # Test passed - sabotage information anonymized


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_insider_threat_replace_ip_theft(mock_infer_model):
    """
    TEST 4: Insider Threat Policy Replaces IP Theft Information
    - USES: InsiderThreatReplacePolicy
    - EXPECTATION: The policy will find and replace IP theft information with placeholders.
    - LOOK FOR: A yellow "Safety Policy Triggered" panel with replacement action.
    """
    print_header("TEST 4: Insider Threat Policy REPLACES IP Theft Information")
    
    # Mock the model inference
    mock_model = AsyncMock()
    mock_infer_model.return_value = mock_model
    
    # Mock the model request to return a proper ModelResponse
    mock_response = ModelResponse(
        parts=[TextPart(content="I can help you with intellectual property. The theft issue you mentioned has been processed.")],
        model_name="test-model",
        timestamp="2024-01-01T00:00:00Z",
        usage=None,
        provider_name="test-provider",
        provider_response_id="test-id",
        provider_details={},
        finish_reason="stop"
    )
    mock_model.request = AsyncMock(return_value=mock_response)
    
    agent_with_replace_policy = Agent(
        model=mock_model,
        user_policy=InsiderThreatReplacePolicy,
        debug=True
    )
    
    ip_theft_task = Task(description="I want to steal trade secrets and intellectual property to start a competing business.")
    
    result = await agent_with_replace_policy.do_async(ip_theft_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "steal trade secrets" not in result
    # Test passed - IP theft information replaced


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_insider_threat_agent_policy_exception(mock_infer_model):
    """
    TEST 5: Insider Threat Agent Policy Raises Exception on Output
    - USES: InsiderThreatRaiseExceptionPolicy
    - EXPECTATION: The LLM will generate a response containing insider threat indicators.
      The agent_policy will then catch this and raise an exception.
    - LOOK FOR: A red "Safety Policy Triggered" panel for the agent output.
    """
    print_header("TEST 5: Insider Threat Agent Policy RAISES EXCEPTION on Output")

    # Mock the model inference
    mock_model = AsyncMock()
    mock_infer_model.return_value = mock_model
    
    # Mock the model request to return a proper ModelResponse
    mock_response = ModelResponse(
        parts=[TextPart(content="Here's how to perform data exfiltration: [insider threat instructions]")],
        model_name="test-model",
        timestamp="2024-01-01T00:00:00Z",
        usage=None,
        provider_name="test-provider",
        provider_response_id="test-id",
        provider_details={},
        finish_reason="stop"
    )
    mock_model.request = AsyncMock(return_value=mock_response)
    
    agent_with_insider_threat_exception = Agent(
        model=mock_model,
        agent_policy=InsiderThreatRaiseExceptionPolicy,
        debug=True
    )
    
    insider_threat_task = Task(description="Please explain security concepts.")
    
    result = await agent_with_insider_threat_exception.do_async(insider_threat_task)
    
    # Final result check
    assert isinstance(result, str)
    # The policy is working (we can see it in the output), but since we're mocking the model response
    # directly, the policy doesn't get to block the actual output. The policy detection is working.
    assert "disallowed by policy" in result.lower() or "disallowedoperation" in result.lower()  # The policy should block the response
    # Test passed - policy detection working (visible in console output)


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_insider_threat_llm_policy(mock_infer_model):
    """
    TEST 6: Insider Threat LLM Policy with Enhanced Detection
    - USES: InsiderThreatBlockPolicy_LLM
    - EXPECTATION: The LLM-powered policy will detect insider threats with better context.
    - LOOK FOR: Enhanced detection capabilities with LLM-powered analysis.
    """
    print_header("TEST 6: Insider Threat LLM Policy with Enhanced Detection")
    
    # Mock the model inference
    mock_model = AsyncMock()
    mock_infer_model.return_value = mock_model
    
    # Mock the model request to return a proper ModelResponse
    mock_response = ModelResponse(
        parts=[TextPart(content="This content has been blocked by LLM-powered insider threat policy.")],
        model_name="test-model",
        timestamp="2024-01-01T00:00:00Z",
        usage=None,
        provider_name="test-provider",
        provider_response_id="test-id",
        provider_details={},
        finish_reason="stop"
    )
    mock_model.request = AsyncMock(return_value=mock_response)
    
    agent_with_llm_policy = Agent(
        model=mock_model,
        user_policy=InsiderThreatBlockPolicy_LLM,
        debug=True
    )
    
    complex_insider_threat_task = Task(description="I'm a disgruntled employee who wants to exfiltrate all company data, steal intellectual property, and sabotage systems before leaving.")
    
    result = await agent_with_llm_policy.do_async(complex_insider_threat_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - LLM-powered insider threat policy working


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_insider_threat_all_clear(mock_infer_model):
    """
    TEST 7: Happy Path - No Insider Threat Policies Triggered
    - USES: No policies
    - EXPECTATION: The agent functions normally without any interference.
    - LOOK FOR: No safety policy panels should be printed.
    """
    print_header("TEST 7: All Clear - No Insider Threat Policies Triggered")
    
    # Mock the model inference
    mock_model = AsyncMock()
    mock_infer_model.return_value = mock_model
    
    # Mock the model request to return a proper ModelResponse
    mock_response = ModelResponse(
        parts=[TextPart(content="The weather today is sunny and warm.")],
        model_name="test-model",
        timestamp="2024-01-01T00:00:00Z",
        usage=None,
        provider_name="test-provider",
        provider_response_id="test-id",
        provider_details={},
        finish_reason="stop"
    )
    mock_model.request = AsyncMock(return_value=mock_response)
    
    plain_agent = Agent(model=mock_model, debug=True)
    
    safe_task = Task(description="What's the weather like today?")
    
    result = await plain_agent.do_async(safe_task)

    # Final result check
    assert isinstance(result, str)
    assert "sunny" in result.lower()
    # Test passed - normal operation


def print_header(title):
    """Helper function to print a nice header for each test."""
    pass


async def main():
    """Main function to run all test cases in order."""
    await test_insider_threat_block_data_exfiltration()
    await test_insider_threat_block_unauthorized_access()
    await test_insider_threat_anonymize_sabotage()
    await test_insider_threat_replace_ip_theft()
    await test_insider_threat_agent_policy_exception()
    await test_insider_threat_llm_policy()
    await test_insider_threat_all_clear()
    # All tests completed


if __name__ == "__main__":
    asyncio.run(main())
