import unittest

import analogai.foundation.extensions.belief_query.tests.env_setup
from analogai.foundation.common.dtos.notion_expression import NotionExpression
from analogai.foundation.core.value_objects.modifier import Modifier
from analogai.foundation.core.value_objects.attribute import Attribute
from analogai.foundation.core.value_objects.time import Time
from analogai.foundation.core.value_objects.relation import Relation
from analogai.foundation.common.enums.deontic_modality import DeonticModality
from analogai.foundation.common.enums.criterion import Criterion
from analogai.foundation.setup.factories import BeliefQueryServiceFactory
from analogai.foundation.tests.functional.services.setup import ServicesTestConfig


def create_learning(
    learnings_service,
    user_id,
    agent_id,
    subject_notion,
    result_criteria,
    modifier=None,
    probability=None,
    validity=None,
    statement_service=None,
    notion_service=None,
    belief_service=None,
):
    relation = result_criteria.get(Criterion.Relation)
    object_name = result_criteria.get(Criterion.ObjectCaption)
    if not relation or not object_name:
        raise ValueError("Belief result requires Relation and ObjectCaption")
    if not isinstance(relation, Relation):
        raise TypeError("Criterion.Relation must be a Relation instance")
    if statement_service is None or notion_service is None:
        raise ValueError("statement_service and notion_service are required to create learning statements")
    if belief_service is None:
        belief_service = learnings_service._belief_service
    attributes = None
    attr_values = result_criteria.get(Criterion.AttributeValue)
    if attr_values:
        attributes = []
        for item in attr_values:
            if isinstance(item, Attribute):
                attributes.append(item)
                continue
            if isinstance(item, dict):
                tag = item.get("tag") or item.get("kind")
                value = item.get("value") or item.get("min_value") or item.get("max_value")
                unit = item.get("unit")
                if tag is not None and value is not None:
                    attributes.append(
                        Attribute(tag=str(tag), value=value, unit=str(unit) if unit else None)
                    )
    complement_notion = notion_service.find_or_create(
        user_id,
        agent_id,
        NotionExpression(caption=str(object_name), attributes=attributes),
    )
    relation_obj = relation
    prob = probability if probability is not None else result_criteria.get(Criterion.RelationProbability, 1.0)
    valid = validity if validity is not None else 1.0
    final_modifier = modifier
    statement_service.create(
        user_id=user_id,
        agent_id=agent_id,
        subject_notion=subject_notion,
        complement_notion=complement_notion,
        relation=relation_obj,
        probability=prob,
        validity=valid,
        modifier=final_modifier,
    )
    return belief_service.find(
        subject_notion=subject_notion,
        complement_notion=complement_notion,
        relation=relation_obj,
        modifier=final_modifier,
    )


class TestBeliefSearch(unittest.TestCase):
    def setUp(self):
        self.config = ServicesTestConfig()
        self.notion_service = self.config.notion_service
        self.statement_service = self.config.statement_service
        self.belief_service = self.config.belief_service
        self.service = BeliefQueryServiceFactory(
            belief_service=self.belief_service,
        ).get_service()
        self.user_id = self.config.user.id
        self.agent_id = self.config.agent.id
        self.subject_notion = self.notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="socrates")
        )

    def _search(self):
        return self.service.search(self.user_id, self.agent_id).exactModifierMatch()

    def tearDown(self):
        self.config.tear_down()

    def test_search_exists_returns_true_when_learning_matches(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        result = self._search().where(
            Criterion.Relation, Relation.from_string("sells")
        ).exists()
        self.assertTrue(result)

    def test_search_exists_returns_false_when_no_match(self):
        result = self._search().where(
            Criterion.Relation, Relation.from_string("nonexistent")
        ).exists()
        self.assertFalse(result)

    def test_search_list_returns_matching_learnings(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        learnings = self._search().where(
            Criterion.SubjectCaption, "socrates"
        ).list()
        self.assertEqual(len(learnings), 1)
        self.assertEqual(learnings[0].relation, Relation.from_string("sells"))

    def test_search_count_returns_correct_count(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("is"),
                Criterion.ObjectCaption: "human",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        count = self._search().count()
        self.assertEqual(count, 2)

    def test_search_has_time_filters_by_time_existence(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            modifier=Modifier(from_time=Time(year=2024, month=1, day=15)),
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "android",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        with_time = self._search().whereNotNull(
            Criterion.Time
        ).list()
        without_time = self._search().whereNull(
            Criterion.Time
        ).list()
        self.assertEqual(len(with_time), 1)
        self.assertEqual(len(without_time), 1)
        self.assertEqual(with_time[0].complement_notion.caption, "iphone")
        self.assertEqual(without_time[0].complement_notion.caption, "android")

    def test_search_min_time_max_time_filter_by_time_range(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            modifier=Modifier(
                from_time=Time(year=2024, month=3, day=1),
                to_time=Time(year=2024, month=3, day=31),
            ),
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "android",
            },
            modifier=Modifier(from_time=Time(year=2024, month=1, day=1)),
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        in_range = self._search().whereMin(
            Criterion.Time, Time(year=2024, month=2, day=1)
        ).whereMax(Criterion.Time, Time(year=2024, month=4, day=1)).list()
        self.assertEqual(len(in_range), 1)
        self.assertEqual(in_range[0].complement_notion.caption, "iphone")

    def test_search_location_filters_by_modifier_location(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            modifier=Modifier(location="store"),
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "android",
            },
            modifier=Modifier(location="online"),
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        at_store = self._search().where(
            Criterion.Location, "store"
        ).list()
        self.assertEqual(len(at_store), 1)
        self.assertEqual(at_store[0].complement_notion.caption, "iphone")

    def test_search_has_deontic_modality_filters_by_deontic_existence(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            modifier=Modifier(deontic_modality=DeonticModality.OBLIGATORY),
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "android",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        with_deontic = self._search().whereNotNull(
            Criterion.DeonticModality
        ).list()
        without_deontic = self._search().whereNull(
            Criterion.DeonticModality
        ).list()
        self.assertEqual(len(with_deontic), 1)
        self.assertEqual(len(without_deontic), 1)
        self.assertEqual(with_deontic[0].complement_notion.caption, "iphone")
        self.assertEqual(without_deontic[0].complement_notion.caption, "android")

    def test_search_deontic_modality_filters_by_deontic_value(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            modifier=Modifier(deontic_modality=DeonticModality.OBLIGATORY),
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "android",
            },
            modifier=Modifier(deontic_modality=DeonticModality.PERMITTED),
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        obligatory = self._search().where(
            Criterion.DeonticModality, DeonticModality.OBLIGATORY
        ).list()
        permitted_by_str = self._search().where(
            Criterion.DeonticModality, "permitted"
        ).list()
        self.assertEqual(len(obligatory), 1)
        self.assertEqual(len(permitted_by_str), 1)
        self.assertEqual(obligatory[0].complement_notion.caption, "iphone")
        self.assertEqual(permitted_by_str[0].complement_notion.caption, "android")

    def test_search_measure_kind_filters_by_measure_kind(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("buys"),
                Criterion.ObjectCaption: "phone",
                Criterion.AttributeValue: [
                    {"tag": "price", "value": 900, "unit": "usd"}
                ],
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("buys"),
                Criterion.ObjectCaption: "book",
                Criterion.AttributeValue: [
                    {"tag": "weight", "value": 2, "unit": "kg"}
                ],
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        priced = self._search().where(
            Criterion.AttributeTag, "price"
        ).list()
        self.assertEqual(len(priced), 1)
        self.assertEqual(priced[0].complement_notion.caption, "phone")

    def test_search_measure_number_whereMax_filters_by_numeric_measure(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("buys"),
                Criterion.ObjectCaption: "phone",
                Criterion.AttributeValue: [
                    {"tag": "price", "value": 900, "unit": "usd"}
                ],
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("buys"),
                Criterion.ObjectCaption: "case",
                Criterion.AttributeValue: [
                    {"tag": "price", "value": 30, "unit": "usd"}
                ],
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        affordable = self._search().whereMax(
            Criterion.AttributeMaxValue, ("price", 100, "usd")
        ).list()
        self.assertEqual(len(affordable), 1)
        self.assertEqual(affordable[0].complement_notion.caption, "case")

    def test_search_or_where_matches_any_value(self):
        notion_service = self.notion_service
        plato = notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="plato")
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            plato,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "android",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        learnings = self._search().where(
            Criterion.SubjectCaption, "socrates"
        ).orWhere(Criterion.SubjectCaption, "plato").list()
        self.assertEqual(len(learnings), 2)
        subjects = {belief.subject_notion.caption for belief in learnings}
        self.assertEqual(subjects, {"socrates", "plato"})

    def test_search_or_where_with_different_criterion(self):
        notion_service = self.notion_service
        mortal = notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="mortal")
        )
        plato = notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="plato")
        )
        notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="is")
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("is"),
                Criterion.ObjectCaption: "mortal",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            plato,
            {
                Criterion.Relation: Relation.from_string("is"),
                Criterion.ObjectCaption: "socrates",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        learnings = self._search().where(
            Criterion.SubjectCaption, "socrates"
        ).orWhere(Criterion.ObjectCaption, "socrates").list()
        self.assertEqual(len(learnings), 2)
        subjects_and_objects = {(b.subject_notion.caption, b.complement_notion.caption) for b in learnings}
        self.assertEqual(subjects_and_objects, {("socrates", "mortal"), ("plato", "socrates")})

    def test_search_where_group_sets_main_clause(self):
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        learnings = self._search().whereGroup({
            Criterion.SubjectCaption: "socrates",
            Criterion.Relation: Relation.from_string("sells"),
        }).list()
        self.assertEqual(len(learnings), 1)
        self.assertEqual(learnings[0].complement_notion.caption, "iphone")

    def test_search_or_where_group_adds_alternative_clause(self):
        notion_service = self.notion_service
        plato = notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="plato")
        )
        notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="mortal")
        )
        notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="is")
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("is"),
                Criterion.ObjectCaption: "mortal",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            plato,
            {
                Criterion.Relation: Relation.from_string("is"),
                Criterion.ObjectCaption: "socrates",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        learnings = self._search().whereGroup({
            Criterion.SubjectCaption: "socrates",
            Criterion.ObjectCaption: "mortal",
        }).orWhereGroup({
            Criterion.SubjectCaption: "plato",
            Criterion.ObjectCaption: "socrates",
        }).list()
        self.assertEqual(len(learnings), 2)
        subjects_and_objects = {(b.subject_notion.caption, b.complement_notion.caption) for b in learnings}
        self.assertEqual(subjects_and_objects, {("socrates", "mortal"), ("plato", "socrates")})

    def test_search_where_in_matches_any_value_in_list(self):
        notion_service = self.notion_service
        plato = notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="plato")
        )
        aristotle = notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="aristotle")
        )
        notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="sells")
        )
        notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="iphone")
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            plato,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            aristotle,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        learnings = self._search().whereIn(
            Criterion.SubjectCaption, ["socrates", "plato"]
        ).where(Criterion.Relation, Relation.from_string("sells")).list()
        self.assertEqual(len(learnings), 2)
        subjects = {b.subject_notion.caption for b in learnings}
        self.assertEqual(subjects, {"socrates", "plato"})

    def test_search_or_where_in_adds_alternative_with_list(self):
        notion_service = self.notion_service
        plato = notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="plato")
        )
        notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="sells")
        )
        notion_service.find_or_create(
            self.user_id, self.agent_id, NotionExpression(caption="iphone")
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            self.subject_notion,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "iphone",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        create_learning(
            self.service,
            self.user_id,
            self.agent_id,
            plato,
            {
                Criterion.Relation: Relation.from_string("sells"),
                Criterion.ObjectCaption: "android",
            },
            statement_service=self.statement_service,
            notion_service=self.notion_service,
        )
        learnings = self._search().where(
            Criterion.SubjectCaption, "socrates"
        ).orWhereIn(Criterion.SubjectCaption, ["plato"]).list()
        self.assertEqual(len(learnings), 2)
        subjects_and_objects = {(b.subject_notion.caption, b.complement_notion.caption) for b in learnings}
        self.assertEqual(subjects_and_objects, {("socrates", "iphone"), ("plato", "android")})



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