import pytest
from atomic_agents.core import PolicyEngine
from atomic_agents.core.policies import PolicyResult


class TestPolicyEngine:
    def test_cost_limit_enforcement(self):
        policies = {"maxCost": 50.0}
        engine = PolicyEngine(policies, mode="enforce")
        
        result = engine.pre_check("llm_call", {"current_cost": 60.0})
        assert result.allowed is False
        assert "Cost limit" in result.reason

    def test_cost_limit_allowed(self):
        policies = {"maxCost": 50.0}
        engine = PolicyEngine(policies, mode="enforce")
        
        result = engine.pre_check("llm_call", {"current_cost": 30.0})
        assert result.allowed is True

    def test_token_limit_enforcement(self):
        policies = {"maxTokens": 1000}
        engine = PolicyEngine(policies, mode="enforce")
        
        result = engine.pre_check("llm_call", {"current_tokens": 1500})
        assert result.allowed is False
        assert "Token limit" in result.reason

    def test_tool_blocking(self):
        policies = {"blockedTools": ["dangerous_tool", "delete_file"]}
        engine = PolicyEngine(policies, mode="enforce")
        
        result = engine.check_tool("dangerous_tool")
        assert result.allowed is False
        assert "blocked" in result.reason.lower()

    def test_tool_allowed(self):
        policies = {"blockedTools": ["dangerous_tool"]}
        engine = PolicyEngine(policies, mode="enforce")
        
        result = engine.check_tool("safe_tool")
        assert result.allowed is True

    def test_warning_mode(self):
        policies = {"maxCost": 50.0}
        engine = PolicyEngine(policies, mode="warn")
        
        result = engine.pre_check("llm_call", {"current_cost": 60.0})
        # In warn mode, should allow but with warnings
        assert result.allowed is True
        assert len(result.warnings) > 0

    def test_empty_policies(self):
        engine = PolicyEngine({}, mode="enforce")
        result = engine.pre_check("llm_call", {})
        assert result.allowed is True

