import unittest
from unittest.mock import Mock

from analogai.foundation.modules.categorical_deduction.categorical_deduction_service import CategoricalDeductionService
from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.direct_statement.direct_statement import DirectStatement
from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.core.value_objects.relation import Relation


class FakeBeliefService:
    def __init__(self, beliefs):
        self._beliefs = list(beliefs)
        self.called_with = None

    def get_all_for_agent(self, agent_id: str):
        self.called_with = agent_id
        return list(self._beliefs)


class TestCategoricalDeductionService(unittest.TestCase):
    def setUp(self):
        self.agent_id = "agent-1"
        self.user_id = "user-1"

    def _build_belief(self, subject_id: str, object_id: str) -> Belief:
        subject = Notion(user_id=self.user_id, agent_id=self.agent_id, caption=subject_id, id=subject_id)
        complement_notion = Notion(user_id=self.user_id, agent_id=self.agent_id, caption=object_id, id=object_id)
        statement = DirectStatement(
            user_id=self.user_id,
            agent_id=self.agent_id,
            subject_notion=subject,
            complement_notion=complement_notion,
            relation=Relation.from_string("is"),
            probability=1.0,
            validity=1.0,
        )
        belief = Belief(agent_id=self.agent_id, id=f"belief-{subject_id}-{object_id}")
        belief.add_statement(statement)
        return belief

    def test_find_deductions_by_agent_id_uses_belief_service_and_paginates(self):
        belief_a = self._build_belief("a", "b")
        belief_b = self._build_belief("b", "c")
        belief_c = self._build_belief("c", "d")
        belief_service = FakeBeliefService([belief_a, belief_b, belief_c])
        service = CategoricalDeductionService(belief_service=belief_service)

        captured = {}

        def fake_resolve_deductions(*, beliefs_by_key, modifier, max_depth):
            captured["beliefs_by_key"] = beliefs_by_key
            captured["modifier"] = modifier
            captured["max_depth"] = max_depth
            return []

        service.resolve_deductions = fake_resolve_deductions

        service.find_deductions_by_agent_id(
            agent_id=self.agent_id,
            limit=1,
            offset=1,
            max_depth=3,
        )

        self.assertEqual(belief_service.called_with, self.agent_id)
        expected_key = service._belief_key_from_values(
            agent_id=belief_b.agent_id,
            subject_id=belief_b.subject_notion.id,
            object_id=belief_b.complement_notion.id,
            relation=belief_b.relation,
            modifier=belief_b.modifier,
        )
        self.assertEqual(set(captured["beliefs_by_key"].keys()), {expected_key})
        self.assertEqual(captured["max_depth"], 3)

    def test_find_deductions_by_agent_id_requires_belief_service(self):
        with self.assertRaises(ValueError):
            CategoricalDeductionService(belief_service=None)


if __name__ == "__main__":
    unittest.main()
