"""Tests for DeepSigmaMemory."""
import pytest

from langchain_deepsigma import DeepSigmaMemory
from coherence_ops import MemoryGraph


class TestDeepSigmaMemory:
    def test_init_defaults(self):
        mem = DeepSigmaMemory()
        assert mem.memory_key == "decision_history"
        assert mem.top_k == 5
        assert isinstance(mem.mg, MemoryGraph)

    def test_memory_variables_property(self):
        mem = DeepSigmaMemory()
        assert mem.memory_variables == ["decision_history"]

    def test_empty_load_returns_empty_string(self):
        mem = DeepSigmaMemory()
        result = mem.load_memory_variables({"input": "hello"})
        assert result == {"decision_history": ""}

    def test_save_context_stores_episode(self):
        mem = DeepSigmaMemory()
        mem.save_context(
            inputs={"input": "What is the capital of France?"},
            outputs={"output": "Paris"},
        )
        assert mem.mg.node_count >= 1  # episode node added

    def test_load_after_save_returns_history(self):
        mem = DeepSigmaMemory()
        mem.save_context(
            inputs={"input": "What is 2+2?"},
            outputs={"output": "4"},
        )
        result = mem.load_memory_variables({"input": "math"})
        history = result["decision_history"]
        assert "Decision History" in history
        assert "What is 2+2?" in history
        assert "4" in history

    def test_top_k_limits_output(self):
        mem = DeepSigmaMemory(top_k=2)
        for i in range(5):
            mem.save_context(
                inputs={"input": f"question {i}"},
                outputs={"output": f"answer {i}"},
            )
        result = mem.load_memory_variables({"input": "something"})
        history = result["decision_history"]
        # Only 2 most recent should appear (top_k=2)
        assert history.count("•") == 2

    def test_clear_resets_memory(self):
        mem = DeepSigmaMemory()
        mem.save_context(inputs={"input": "hello"}, outputs={"output": "hi"})
        assert mem.mg.node_count >= 1
        mem.clear()
        assert mem.mg.node_count == 0
        result = mem.load_memory_variables({"input": "test"})
        assert result["decision_history"] == ""

    def test_custom_memory_key(self):
        mem = DeepSigmaMemory(memory_key="past_decisions")
        assert mem.memory_variables == ["past_decisions"]
        mem.save_context(inputs={"input": "test"}, outputs={"output": "result"})
        result = mem.load_memory_variables({})
        assert "past_decisions" in result

    def test_shared_mg_with_callback(self):
        from langchain_deepsigma import DeepSigmaCallbackHandler
        import uuid

        shared_mg = MemoryGraph()
        mem = DeepSigmaMemory(mg=shared_mg)
        handler = DeepSigmaCallbackHandler()
        handler.mg = shared_mg  # Share the same graph

        # Add episode via memory
        mem.save_context(inputs={"input": "q"}, outputs={"output": "a"})
        assert shared_mg.node_count >= 1
