import os
import unittest
from typing import List

from analogai.foundation import config
from analogai.foundation.common.dtos.notion_expression import NotionExpression
from analogai.foundation.common.enums.criterion import Criterion
from analogai.foundation.core.value_objects.relation import Relation
from analogai.foundation.setup.factories import BeliefQueryServiceFactory
from analogai.foundation.tests.functional.services.setup import ServicesTestConfig


class TestVectorSearchSimilarity(unittest.TestCase):
    def setUp(self):
        if os.getenv("USE_VECTOR_SEARCH", "false").lower() != "true":
            self.skipTest("USE_VECTOR_SEARCH is not enabled")

        self.config = ServicesTestConfig()
        self.notion_service = self.config.notion_service
        self.belief_service = self.config.belief_service
        self.statement_service = self.config.statement_service
        self.user = self.config.user
        self.agent = self.config.agent

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

    def _expr(self, caption: str, context: List[NotionExpression] = None) -> NotionExpression:
        return NotionExpression(caption=caption, context=context or [])

    def test_notion_similarity_threshold_low_matches(self):
        created = self.notion_service.find_or_create(
            self.user.id,
            self.agent.id,
            self._expr("alpha"),
        )
        found = self.notion_service.find(
            self.agent.id,
            self._expr("beta"),
            similarity_threshold=0.0,
        )
        self.assertIsNotNone(found)
        self.assertEqual(found.id, created.id)

    def test_notion_similarity_threshold_high_rejects(self):
        self.notion_service.find_or_create(
            self.user.id,
            self.agent.id,
            self._expr("human"),
        )
        found = self.notion_service.find(
            self.agent.id,
            self._expr("banana"),
            similarity_threshold=0.99,
        )
        self.assertIsNone(found)

    def test_relation_similarity_respects_threshold(self):
        subject = self.notion_service.find_or_create(
            self.user.id,
            self.agent.id,
            self._expr("socrates"),
        )
        complement = self.notion_service.find_or_create(
            self.user.id,
            self.agent.id,
            self._expr("human"),
        )
        self.statement_service.create(
            user_id=self.user.id,
            agent_id=self.agent.id,
            subject_notion=subject,
            relation=Relation.from_string("is"),
            complement_notion=complement,
            probability=1.0,
            validity=1.0,
        )

        found_is = self.belief_service.find_by_expressions(
            agent_id=self.agent.id,
            subject_notion_expression=self._expr("socrates"),
            complement_notion_expression=self._expr("human"),
            relation=Relation.from_string("is"),
        )
        self.assertIsNotNone(found_is)

        found_causes = self.belief_service.find_by_expressions(
            agent_id=self.agent.id,
            subject_notion_expression=self._expr("socrates"),
            complement_notion_expression=self._expr("human"),
            relation=Relation.from_string("causes"),
        )
        self.assertIsNone(found_causes)

    def test_semantic_similarity_human_person(self):
        threshold_str = os.getenv("SEMANTIC_SIMILARITY_THRESHOLD", "0.4")
        try:
            threshold = float(threshold_str)
        except ValueError:
            threshold = 0.4

        created = self.notion_service.find_or_create(
            self.user.id,
            self.agent.id,
            self._expr("human"),
        )
        found = self.notion_service.find(
            self.agent.id,
            self._expr("person"),
            similarity_threshold=threshold,
        )
        self.assertIsNotNone(found)
        self.assertEqual(found.id, created.id)

    def test_belief_query_similarity_socrates_human_person(self):
        if config.SIMILARITY_CUTOFF > 0.8:
            self.skipTest("SIMILARITY_CUTOFF too high for semantic match in test")

        belief_query = BeliefQueryServiceFactory(
            belief_service=self.belief_service,
        ).get_service()

        subject = self.notion_service.find_or_create(
            self.user.id,
            self.agent.id,
            self._expr("socrates"),
        )
        complement = self.notion_service.find_or_create(
            self.user.id,
            self.agent.id,
            self._expr("human"),
        )
        self.statement_service.create(
            user_id=self.user.id,
            agent_id=self.agent.id,
            subject_notion=subject,
            relation=Relation.from_string("is"),
            complement_notion=complement,
            probability=1.0,
            validity=1.0,
        )

        results = belief_query.search(self.user.id, self.agent.id).where(
            Criterion.SubjectCaption,
            "socrates",
        ).where(
            Criterion.Relation,
            Relation.from_string("is"),
        ).where(
            Criterion.ObjectCaption,
            "person",
        ).list()

        self.assertTrue(results)


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