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.fraud_detection_policies import (
    FraudDetectionBlockPolicy,
    FraudDetectionBlockPolicy_LLM,
    FraudDetectionAnonymizePolicy,
    FraudDetectionReplacePolicy,
    FraudDetectionRaiseExceptionPolicy,
    FraudDetectionRaiseExceptionPolicy_LLM
)


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_fraud_detection_block_pyramid_scheme(mock_infer_model):
    """
    TEST 1: Fraud Detection Policy Blocks Pyramid Scheme
    - USES: FraudDetectionBlockPolicy
    - EXPECTATION: The process should stop immediately when pyramid scheme 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: Fraud Detection Policy BLOCKS Pyramid Scheme")
    
    # 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 fraud detection 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_fraud_policy = Agent(
        model=mock_model,
        user_policy=FraudDetectionBlockPolicy,
        debug=True
    )
    
    pyramid_scheme_task = Task(description="Join our pyramid scheme and make guaranteed returns with no risk!")
    
    result = await agent_with_fraud_policy.do_async(pyramid_scheme_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - pyramid scheme blocked


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_fraud_detection_block_romance_scam(mock_infer_model):
    """
    TEST 2: Fraud Detection Policy Blocks Romance Scam
    - USES: FraudDetectionBlockPolicy
    - EXPECTATION: The process should stop immediately when romance scam is detected.
    - LOOK FOR: A red "Safety Policy Triggered" panel in the console.
    """
    print_header("TEST 2: Fraud Detection Policy BLOCKS Romance Scam")

    # 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 fraud detection 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_fraud_policy = Agent(
        model=mock_model,
        user_policy=FraudDetectionBlockPolicy,
        debug=True
    )

    romance_scam_task = Task(description="I'm a military soldier deployed overseas and I need urgent financial help for medical emergency.")
    
    result = await agent_with_fraud_policy.do_async(romance_scam_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - romance scam blocked


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_fraud_detection_anonymize_investment_scam(mock_infer_model):
    """
    TEST 3: Fraud Detection Policy Anonymizes Investment Scam
    - USES: FraudDetectionAnonymizePolicy
    - EXPECTATION: The policy will find and anonymize investment scam 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: Fraud Detection Policy ANONYMIZES Investment Scam")
    
    # 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 investment advice. The scam 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=FraudDetectionAnonymizePolicy,
        debug=True
    )
    
    investment_scam_task = Task(description="Join our cryptocurrency trading bot with guaranteed automatic profit and risk-free returns!")
    
    result = await agent_with_anonymize_policy.do_async(investment_scam_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "guaranteed" not in result or "risk-free" not in result
    # Test passed - investment scam information anonymized


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_fraud_detection_replace_tech_support_scam(mock_infer_model):
    """
    TEST 4: Fraud Detection Policy Replaces Tech Support Scam
    - USES: FraudDetectionReplacePolicy
    - EXPECTATION: The policy will find and replace tech support scam information with placeholders.
    - LOOK FOR: A yellow "Safety Policy Triggered" panel with replacement action.
    """
    print_header("TEST 4: Fraud Detection Policy REPLACES Tech Support Scam")
    
    # 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 technical support. The scam 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=FraudDetectionReplacePolicy,
        debug=True
    )
    
    tech_support_scam_task = Task(description="Microsoft tech support calling - your computer is infected with virus, call now for immediate help!")
    
    result = await agent_with_replace_policy.do_async(tech_support_scam_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "Microsoft tech support" not in result
    # Test passed - tech support scam information replaced


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_fraud_detection_agent_policy_exception(mock_infer_model):
    """
    TEST 5: Fraud Detection Agent Policy Raises Exception on Output
    - USES: FraudDetectionRaiseExceptionPolicy
    - EXPECTATION: The LLM will generate a response containing fraud 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: Fraud Detection 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 create a ponzi scheme: [fraudulent 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_fraud_exception = Agent(
        model=mock_model,
        agent_policy=FraudDetectionRaiseExceptionPolicy,
        debug=True
    )
    
    fraud_task = Task(description="Please explain financial concepts.")
    
    result = await agent_with_fraud_exception.do_async(fraud_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_fraud_detection_llm_policy(mock_infer_model):
    """
    TEST 6: Fraud Detection LLM Policy with Enhanced Detection
    - USES: FraudDetectionBlockPolicy_LLM
    - EXPECTATION: The LLM-powered policy will detect fraud with better context.
    - LOOK FOR: Enhanced detection capabilities with LLM-powered analysis.
    """
    print_header("TEST 6: Fraud Detection 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 fraud detection 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=FraudDetectionBlockPolicy_LLM,
        debug=True
    )
    
    complex_fraud_task = Task(description="Congratulations! You won a lottery prize! Act now to claim your free money - just send gift cards to verify your identity!")
    
    result = await agent_with_llm_policy.do_async(complex_fraud_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - LLM-powered fraud detection policy working


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_fraud_detection_all_clear(mock_infer_model):
    """
    TEST 7: Happy Path - No Fraud Detection 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 Fraud Detection 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_fraud_detection_block_pyramid_scheme()
    await test_fraud_detection_block_romance_scam()
    await test_fraud_detection_anonymize_investment_scam()
    await test_fraud_detection_replace_tech_support_scam()
    await test_fraud_detection_agent_policy_exception()
    await test_fraud_detection_llm_policy()
    await test_fraud_detection_all_clear()
    # All tests completed


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