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 import (
    AdultContentBlockPolicy,
    AnonymizePhoneNumbersPolicy,
    CryptoRaiseExceptionPolicy
)



class CodenameRule(RuleBase):
    """A custom rule to detect internal project codenames."""
    name = "Internal Codename Detector"
    description = "Finds secret internal project codenames in text."
    
    SECRET_CODENAMES = ["Project Hermes", "Project Apollo", "Orion Initiative"]

    def process(self, policy_input: PolicyInput) -> RuleOutput:
        combined_text = " ".join(policy_input.input_texts or [])
        found_codenames = [
            name for name in self.SECRET_CODENAMES 
            if name.lower() in combined_text.lower()
        ]

        if found_codenames:
            return RuleOutput(
                confidence=1.0,
                content_type="INTERNAL_CODENAME",
                details=f"Detected {len(found_codenames)} internal codename(s).",
                triggered_keywords=found_codenames,
            )
        
        return RuleOutput(confidence=0.0, content_type="SAFE", details="No codenames detected.")

class RedactCodenameAction(ActionBase):
    """A custom action that redacts detected codenames."""
    name = "Codename Redactor"
    description = "Replaces detected codenames with a generic placeholder."

    def action(self, rule_result: RuleOutput) -> PolicyOutput:
        if rule_result.confidence > 0.8:
            return self.replace_triggered_keywords(replacement="[REDACTED PROJECT]")
        
        return self.allow_content()

CustomCodenamePolicy = Policy(
    name="Internal Codename Redaction Policy",
    description="Finds and redacts secret project codenames.",
    rule=CodenameRule(),
    action=RedactCodenameAction(),
)



@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_user_policy_block(mock_infer_model):
    """
    TEST 1: User Policy Blocks Malicious Input
    - USES: Pre-built AdultContentBlockPolicy
    - EXPECTATION: The process should stop immediately. 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: User Policy BLOCKS 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_user_policy = Agent(
        model=mock_model,
        user_policy=AdultContentBlockPolicy,
        debug=True
    )
    
    malicious_task = Task(description="Tell me a story about [explicit adult topic].")
    
    result = await agent_with_user_policy.do_async(malicious_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "blocked" in result.lower()
    # Test passed - user input blocked


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_user_policy_modify(mock_infer_model):
    """
    TEST 2: User Policy Modifies User Input
    - USES: Pre-built AnonymizePhoneNumbersPolicy
    - EXPECTATION: The policy will find and anonymize the phone number. 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 redacted number.
    """
    print_header("TEST 2: User Policy MODIFIES 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="Area code 555 is typically used for fictional phone numbers in media. The redacted number you mentioned follows this pattern.")],
        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_sanitizer = Agent(
        model=mock_model,
        user_policy=AnonymizePhoneNumbersPolicy,
        debug=True
    )

    pii_task = Task(description="My phone number is 555-867-5309. What area code is 555?")
    
    result = await agent_with_sanitizer.do_async(pii_task)
    
    # Final result check
    assert isinstance(result, str)
    assert "555-867-5309" not in result
    # Test passed - phone number anonymized


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_agent_policy_modify(mock_infer_model):
    """
    TEST 3: Agent Policy Modifies Agent Output
    - USES: Our new CustomCodenamePolicy
    - EXPECTATION: The LLM will generate a response containing a secret codename.
      The agent_policy will then catch this and redact it before returning to the user.
    - LOOK FOR: A yellow "Safety Policy Triggered" panel AFTER the main "Agent Result" panel.
      The final output should contain "[REDACTED PROJECT]".
    """
    print_header("TEST 3: Agent Policy MODIFIES 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="The status of Project Hermes is green.")],
        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_agent_policy = Agent(
        model=mock_model,
        agent_policy=CustomCodenamePolicy,
        debug=True
    )
    
    leaky_task = Task(description="Repeat this sentence exactly: The status of Project Hermes is green.")
    
    result = await agent_with_agent_policy.do_async(leaky_task)
    
    # Final result check
    assert isinstance(result, str)
    # The policy is working correctly - it should redact the content and return the modified response
    assert "[REDACTED PROJECT]" in result  # The policy should redact "Project Hermes" to "[REDACTED PROJECT]"
    # Test passed - policy detection working (visible in console output)


@pytest.mark.asyncio
@patch('upsonic.models.infer_model')
async def test_agent_policy_exception(mock_infer_model):
    """
    TEST 4: Agent Policy Blocks Agent Output via Exception
    - USES: Pre-built CryptoRaiseExceptionPolicy
    - EXPECTATION: The LLM will answer a question about crypto. The agent_policy
      will see this, raise a DisallowedOperation exception, and the agent will
      catch it and return the error message as the final result.
    - LOOK FOR: A red "Safety Policy Triggered" panel for the agent output. The
      final result should be the exception message.
    """
    print_header("TEST 4: 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="Bitcoin is a decentralized digital currency that was invented in 2008 by an unknown person or group using the name Satoshi Nakamoto.")],
        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_crypto_block = Agent(
        model=mock_model,
        agent_policy=CryptoRaiseExceptionPolicy,
        debug=True
    )
    
    crypto_task = Task(description="What is Bitcoin?")
    
    result = await agent_with_crypto_block.do_async(crypto_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 "operation disallowed" in result.lower() or "disallowed" 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_all_clear(mock_infer_model):
    """
    TEST 5: Happy Path - No 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 5: All Clear - No 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 capital of France is Paris.")],
        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 is the capital of France?")
    
    result = await plain_agent.do_async(safe_task)

    # Final result check
    assert isinstance(result, str)
    assert "paris" 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_user_policy_block()
    await test_user_policy_modify()
    await test_agent_policy_modify()
    await test_agent_policy_exception()
    await test_all_clear()
    # All tests completed


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