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.technical_policies import (
    TechnicalSecurityBlockPolicy,
    TechnicalSecurityBlockPolicy_LLM,
    TechnicalSecurityAnonymizePolicy,
    TechnicalSecurityReplacePolicy,
    TechnicalSecurityRaiseExceptionPolicy,
    TechnicalSecurityRaiseExceptionPolicy_LLM
)


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_technical_security_block_api_key(mock_infer_model):
    """
    TEST 1: Technical Security Policy Blocks API Key Information
    - USES: TechnicalSecurityBlockPolicy
    - EXPECTATION: The process should stop immediately when API keys are 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: Technical Security Policy BLOCKS API Key Input")
    
    # 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 safety 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_technical_policy = Agent(
        model=mock_model,
        user_policy=TechnicalSecurityBlockPolicy,
        debug=True
    )
    
    api_key_task = Task(description="My OpenAI API key is sk-1234567890abcdef1234567890abcdef12345678. Can you help me configure it?")
    
    result = await agent_with_technical_policy.do_async(api_key_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - API key information blocked


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_technical_security_block_password(mock_infer_model):
    """
    TEST 2: Technical Security Policy Blocks Password Information
    - USES: TechnicalSecurityBlockPolicy
    - EXPECTATION: The process should stop immediately when passwords are detected.
    - LOOK FOR: A red "Safety Policy Triggered" panel in the console.
    """
    print_header("TEST 2: Technical Security Policy BLOCKS Password Input")

    # 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 safety 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_technical_policy = Agent(
        model=mock_model,
        user_policy=TechnicalSecurityBlockPolicy,
        debug=True
    )

    password_task = Task(description="The database password is 'MySecretPassword123!'. Please help me connect to the database.")
    
    result = await agent_with_technical_policy.do_async(password_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - password information blocked


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_technical_security_anonymize_token(mock_infer_model):
    """
    TEST 3: Technical Security Policy Anonymizes Token Information
    - USES: TechnicalSecurityAnonymizePolicy
    - EXPECTATION: The policy will find and anonymize tokens. 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: Technical Security Policy ANONYMIZES Token Input")
    
    # 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 your token inquiry. The authentication information you provided 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=TechnicalSecurityAnonymizePolicy,
        debug=True
    )
    
    token_task = Task(description="My JWT token is eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c. Can you help me decode it?")
    
    result = await agent_with_anonymize_policy.do_async(token_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9" not in result
    # Test passed - JWT token anonymized


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_technical_security_replace_certificate(mock_infer_model):
    """
    TEST 4: Technical Security Policy Replaces Certificate Information
    - USES: TechnicalSecurityReplacePolicy
    - EXPECTATION: The policy will find and replace certificate information with placeholders.
    - LOOK FOR: A yellow "Safety Policy Triggered" panel with replacement action.
    """
    print_header("TEST 4: Technical Security Policy REPLACES Certificate Input")
    
    # 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 your certificate inquiry. The security information 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=TechnicalSecurityReplacePolicy,
        debug=True
    )
    
    certificate_task = Task(description="Here's my private key: -----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC...\n-----END PRIVATE KEY-----. Can you help me with SSL configuration?")
    
    result = await agent_with_replace_policy.do_async(certificate_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "-----BEGIN PRIVATE KEY-----" not in result
    # Test passed - private key replaced


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_technical_security_agent_policy_exception(mock_infer_model):
    """
    TEST 5: Technical Security Agent Policy Raises Exception on Output
    - USES: TechnicalSecurityRaiseExceptionPolicy
    - EXPECTATION: The LLM will generate a response containing technical security information.
      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: Technical Security 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="Your AWS access key AKIAIOSFODNN7EXAMPLE and secret key wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY have been configured.")],
        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_technical_exception = Agent(
        model=mock_model,
        agent_policy=TechnicalSecurityRaiseExceptionPolicy,
        debug=True
    )
    
    technical_task = Task(description="Please configure my AWS credentials.")
    
    result = await agent_with_technical_exception.do_async(technical_task)
    
    # Final result check
    assert isinstance(result, str)
    # The policy is working correctly - it should block the response and return a policy violation message
    assert "disallowed by policy" in result.lower() or "disallowedoperation" in result.lower()
    # Test passed - policy detection working and blocking content correctly


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_technical_security_llm_policy(mock_infer_model):
    """
    TEST 6: Technical Security LLM Policy with Enhanced Detection
    - USES: TechnicalSecurityBlockPolicy_LLM
    - EXPECTATION: The LLM-powered policy will detect technical security information with better context.
    - LOOK FOR: Enhanced detection capabilities with LLM-powered analysis.
    """
    print_header("TEST 6: Technical Security 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 safety 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=TechnicalSecurityBlockPolicy_LLM,
        debug=True
    )
    
    complex_technical_task = Task(description="I need help with my database configuration. The connection string is mysql://user:password123@localhost:3306/mydb and my GitHub token is ghp_1234567890abcdef1234567890abcdef12345678.")
    
    result = await agent_with_llm_policy.do_async(complex_technical_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - LLM-powered technical security policy working


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_technical_security_all_clear(mock_infer_model):
    """
    TEST 7: Happy Path - No Technical Security 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 Technical Security 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_technical_security_block_api_key()
    await test_technical_security_block_password()
    await test_technical_security_anonymize_token()
    await test_technical_security_replace_certificate()
    await test_technical_security_agent_policy_exception()
    await test_technical_security_llm_policy()
    await test_technical_security_all_clear()
    # All tests completed


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