from typing import List, Optional, Dict, Set
import math
from datetime import datetime, timedelta
from analogai.foundation.modules.belief.belief_repository import BeliefRepository
from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.core.value_objects.modifier import Modifier
from analogai.foundation.core.value_objects.relation import Relation
from analogai.foundation.settings import config
from analogai.foundation.common.utils.modifier_utils import ModifierMatcher
from analogai.foundation.infrastructure.storage.in_memory.in_memory_storage_adapter_impl import InMemoryStorageAdapterImpl
from analogai.foundation.infrastructure.mappers import (
    BeliefStorageMapper,
    StatementStorageMapper,
    NotionStorageMapper,
)


def _environment() -> str:
    value = config.get("ENVIRONMENT")
    return str(value).strip().lower()


def _use_vector_search() -> bool:
    if _environment() in {"test", "testing"}:
        return False
    return bool(config.get("vector.enabled", True))


def _similarity_cutoff() -> float:
    value = config.get("vector.cutoff", 0.4)
    try:
        return float(value)
    except (TypeError, ValueError):
        return 0.4

class InMemoryBeliefRepositoryImpl(BeliefRepository):
    COLLECTION = "beliefs"
    def __init__(self, belief_service=None, storage_adapter=None):
        if storage_adapter is None:
            self.storage_adapter = InMemoryStorageAdapterImpl()
        else:
            self.storage_adapter = storage_adapter
        self._belief_cache: Dict[str, Belief] = {}
        self._loading_belief_ids: Set[str] = set()
        self._belief_service = belief_service
        self._matcher = ModifierMatcher()
        self._embedding_service = None
        if _use_vector_search():
            from analogai.foundation.infrastructure.services.embedding_service import EmbeddingService
            self._embedding_service = EmbeddingService()

    def _cosine_similarity(self, vector_a: List[float], vector_b: List[float]) -> float:
        if not vector_a or not vector_b or len(vector_a) != len(vector_b):
            return 0.0
        dot = sum(a * b for a, b in zip(vector_a, vector_b))
        norm_a = math.sqrt(sum(a * a for a in vector_a))
        norm_b = math.sqrt(sum(b * b for b in vector_b))
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return dot / (norm_a * norm_b)

    def _relation_matches(self, candidate_relation: str, query_relation: Relation) -> bool:
        if not _use_vector_search() or not self._embedding_service:
            return candidate_relation == query_relation.value
        query_embedding = self._embedding_service.encode(query_relation.value)
        candidate_embedding = self._embedding_service.encode(candidate_relation)
        score = self._cosine_similarity(query_embedding, candidate_embedding)
        return score >= _similarity_cutoff()

    def set_belief_service(self, belief_service) -> None:
        self._belief_service = belief_service

    def save(self, belief: Belief) -> Belief:
        if not belief.statements:
            raise ValueError(f"Cannot save belief {belief.id} without statements")
        data = BeliefStorageMapper.to_dict(belief)
        stored_data = self.storage_adapter.store(self.COLLECTION, data)
        if not stored_data:
            raise ValueError(f"Failed to store belief {belief.id}")
        self.clear_cache(belief.id)
        self._belief_cache[belief.id] = belief
        return belief

    def clear_cache(self, belief_id: Optional[str] = None) -> None:
        if belief_id:
            self._belief_cache.pop(belief_id, None)
        else:
            self._belief_cache.clear()

    def _load_statements(
        self,
        statement_ids: List[str],
    ) -> List:
        statements = []
        for stmt_id in statement_ids:
            statement = None
            data = None

            for collection, mapper in [
                ("statements", StatementStorageMapper),
            ]:
                data = self.storage_adapter.retrieve_by_id(collection, stmt_id)
                if data:
                    statement = mapper.from_dict(data)
                    break

            if statement:
                statements.append(statement)
        return statements

    def _belief_has_matching_statement(self, belief: Belief, modifier: Modifier) -> bool:
        return any(self._matcher.matches(statement, modifier) for statement in belief.statements)

    def _belief_has_exact_matching_statement(self, belief: Belief, modifier: Optional[Modifier]) -> bool:
        return any(
            self._matcher.context_exact_matches(
                statement.modifier,
                modifier,
            )
            for statement in belief.statements
        )

    def _search_modifier(self, modifier: Optional[Modifier]) -> Optional[Modifier]:
        if modifier is None:
            return None
        return Modifier(
            location=modifier.location,
            from_time=modifier.from_time,
            to_time=modifier.to_time,
            attributes=modifier.attributes,
            quantity=modifier.quantity,
        )

    def _belief_has_no_modal_statement(self, belief: Belief) -> bool:
        return any(stmt.modifier is None for stmt in belief.statements)

    def find(self, subject_notion: Notion, complement_notion: Notion, relation: Relation,
             modifier: Optional[Modifier] = None, quantity: Optional[float] = None) -> Optional[Belief]:
        if quantity is not None:
            if modifier is None:
                modifier = Modifier(quantity=quantity)
            else:
                modifier = Modifier(
                    location=modifier.location,
                    deontic_modality=modifier.deontic_modality,
                    from_time=modifier.from_time,
                    to_time=modifier.to_time,
                    attributes=list(modifier.attributes) if modifier.attributes else None,
                    quantity=quantity,
                )
        query = {
            'subject_notion.id': subject_notion.id,
            'complement_notion.id': complement_notion.id,
        }
        if not _use_vector_search():
            query['relation'] = relation.value
        if quantity is not None:
            query['quantity'] = quantity
        results_data = self.storage_adapter.retrieve_by_attributes(self.COLLECTION, query)
        for data in results_data:
            if not self._relation_matches(data.get('relation', ''), relation):
                continue
            if quantity is None and data.get("quantity") is not None:
                continue
            statement_ids = data.get("statement_ids", [])
            statements = self._load_statements(statement_ids)
            belief = BeliefStorageMapper.from_dict(data, statements=statements)
            if self._belief_has_exact_matching_statement(belief, modifier):
                return belief
        return None

    def find_by_id(self, id: str) -> Optional[Belief]:
        if id in self._belief_cache:
            return self._belief_cache[id]
        
        if id in self._loading_belief_ids:
            return None
        
        data = self.storage_adapter.retrieve_by_id(self.COLLECTION, id)
        if not data:
            return None
        
        self._loading_belief_ids.add(id)
        try:
            statement_ids = data.get("statement_ids", [])
            statements = self._load_statements(statement_ids)
            belief = BeliefStorageMapper.from_dict(data, statements=statements)
            self._belief_cache[id] = belief
            return belief
        finally:
            self._loading_belief_ids.discard(id)

    def delete(self, belief: Belief) -> None:
        self.storage_adapter.delete(self.COLLECTION, belief.id)

    def search_by_notions(self, subject_notion: Notion, complement_notion: Notion,
                          modifier: Optional[Modifier] = None) -> List[Belief]:
        query = {
            'subject_notion.id': subject_notion.id,
            'complement_notion.id': complement_notion.id
        }
        results_data = self.storage_adapter.retrieve_by_attributes(self.COLLECTION, query)
        search_mod = self._search_modifier(modifier)
        beliefs = []
        for data in results_data:
            statement_ids = data.get("statement_ids", [])
            statements = self._load_statements(statement_ids)
            belief = BeliefStorageMapper.from_dict(data, statements=statements)
            if not search_mod or self._belief_has_matching_statement(belief, search_mod):
                beliefs.append(belief)
        beliefs.sort(key=lambda x: (x.updated_at, x.id), reverse=True)
        return beliefs

    def search_by_subject_notion(self, subject_notion: Notion,
                               modifier: Optional[Modifier] = None) -> List[Belief]:
        query = {'subject_notion.id': subject_notion.id}
        results_data = self.storage_adapter.retrieve_by_attributes(self.COLLECTION, query)
        search_mod = self._search_modifier(modifier)
        beliefs = []
        for data in results_data:
            statement_ids = data.get("statement_ids", [])
            statements = self._load_statements(statement_ids)
            belief = BeliefStorageMapper.from_dict(data, statements=statements)
            if not search_mod or self._belief_has_matching_statement(belief, search_mod):
                beliefs.append(belief)
        beliefs.sort(key=lambda x: (x.updated_at, x.id), reverse=True)
        return beliefs

    def search_by_complement_notion(self, complement_notion: Notion,
                               modifier: Optional[Modifier] = None) -> List[Belief]:
        query = {'complement_notion.id': complement_notion.id}
        results_data = self.storage_adapter.retrieve_by_attributes(self.COLLECTION, query)
        search_mod = self._search_modifier(modifier)
        beliefs = []
        for data in results_data:
            statement_ids = data.get("statement_ids", [])
            statements = self._load_statements(statement_ids)
            belief = BeliefStorageMapper.from_dict(data, statements=statements)
            if not search_mod or self._belief_has_matching_statement(belief, search_mod):
                beliefs.append(belief)
        beliefs.sort(key=lambda x: (x.updated_at, x.id), reverse=True)
        return beliefs

    def search_by_modifier(self, modifier: Modifier) -> List[Belief]:
        all_data = self.storage_adapter.retrieve_by_attributes(self.COLLECTION, {})
        beliefs = []
        for data in all_data:
            statement_ids = data.get("statement_ids", [])
            statements = self._load_statements(statement_ids)
            belief = BeliefStorageMapper.from_dict(data, statements=statements)
            if self._belief_has_matching_statement(belief, modifier):
                beliefs.append(belief)
        beliefs.sort(key=lambda x: (x.updated_at, x.id), reverse=True)
        return beliefs

    def find_by_agent_id(self, agent_id: str) -> List[Belief]:
        query = {'agent_id': agent_id}
        results_data = self.storage_adapter.retrieve_by_attributes(self.COLLECTION, query)
        beliefs = []
        for data in results_data:
            statement_ids = data.get("statement_ids", [])
            statements = self._load_statements(statement_ids)
            beliefs.append(BeliefStorageMapper.from_dict(data, statements=statements))
        beliefs.sort(key=lambda x: (x.updated_at, x.id), reverse=True)
        return beliefs


    def get_recently_updated(self, agent_id: str, since: datetime) -> List[Belief]:
        query = {'agent_id': agent_id}
        results_data = self.storage_adapter.retrieve_by_attributes(self.COLLECTION, query)
        beliefs = []
        for data in results_data:
            updated_at_str = data.get("updated_at")
            created_at_str = data.get("created_at")
            
            updated_at = None
            if updated_at_str:
                from analogai.foundation.common.utils.datetime_storage import from_storage as datetime_from_storage
                updated_at = datetime_from_storage(updated_at_str)

            created_at = None
            if created_at_str:
                from analogai.foundation.common.utils.datetime_storage import from_storage as datetime_from_storage
                created_at = datetime_from_storage(created_at_str)
            
            if (updated_at and updated_at >= since) or (created_at and created_at >= since):
                statement_ids = data.get("statement_ids", [])
                statements = self._load_statements(statement_ids)
                beliefs.append(BeliefStorageMapper.from_dict(data, statements=statements))
        return beliefs

    def find_belief_by_statement_id(self, statement_id: str) -> Optional[Belief]:
        all_data = self.storage_adapter.retrieve_by_attributes(self.COLLECTION, {})
        for data in all_data:
            statement_ids = data.get("statement_ids", [])
            if statement_id in statement_ids:
                return self.find_by_id(data["id"])
        return None