import pytest
from uuid import uuid4
from atomic_agents.integrations.langchain import AtomicOptimizer
from langchain_core.outputs import LLMResult, Generation
from langchain_core.agents import AgentAction


class TestAtomicOptimizer:
    def test_initialization(self):
        optimizer = AtomicOptimizer(
            compression=True,
            policies={"blockedTools": ["test"]},
            ledger=True
        )
        assert optimizer.compressor is not None
        assert optimizer.policy_engine is not None
        assert optimizer.ledger is not None

    def test_initialization_no_compression(self):
        optimizer = AtomicOptimizer(compression=False)
        assert optimizer.compressor is None

    def test_initialization_no_ledger(self):
        optimizer = AtomicOptimizer(ledger=False)
        assert optimizer.ledger is None

    def test_on_llm_start_compression(self):
        optimizer = AtomicOptimizer(compression=True, ledger=False)
        run_id = uuid4()
        prompts = ["Please help me. The goal is to summarize."]
        original_size = len(prompts[0])
        
        optimizer.on_llm_start({}, prompts, run_id=run_id)
        
        # Prompt should be compressed
        assert len(prompts[0]) < original_size or "→" in prompts[0] or "⊢" in prompts[0]

    def test_on_llm_start_ledger(self):
        optimizer = AtomicOptimizer(compression=False, ledger=True)
        run_id = uuid4()
        prompts = ["test prompt"]
        
        optimizer.on_llm_start({}, prompts, run_id=run_id)
        
        assert len(optimizer.ledger.entries) > 0
        assert optimizer.ledger.entries[-1]["data"]["event"] == "llm_start"

    def test_on_llm_end_ledger(self):
        optimizer = AtomicOptimizer(compression=False, ledger=True)
        run_id = uuid4()
        prompts = ["test"]
        optimizer.on_llm_start({}, prompts, run_id=run_id)
        
        llm_result = LLMResult(generations=[[Generation(text="response")]])
        optimizer.on_llm_end(llm_result, run_id=run_id)
        
        events = [e["data"]["event"] for e in optimizer.ledger.entries]
        assert "llm_end" in events

    def test_policy_violation_cost(self):
        optimizer = AtomicOptimizer(
            compression=False,
            policies={"maxCost": 10.0},
            ledger=False
        )
        optimizer.current_cost = 15.0
        run_id = uuid4()
        prompts = ["test"]
        
        with pytest.raises(Exception, match="Policy Violation"):
            optimizer.on_llm_start({}, prompts, run_id=run_id)

    def test_tool_blocking(self):
        optimizer = AtomicOptimizer(
            compression=False,
            policies={"blockedTools": ["dangerous_tool"]},
            ledger=False
        )
        run_id = uuid4()
        action = AgentAction(tool="dangerous_tool", tool_input="", log="")
        
        with pytest.raises(Exception, match="Tool Blocked"):
            optimizer.on_agent_action(action, run_id=run_id)


