from typing import List, Optional
import math

from analogai.foundation.settings import config
from analogai.foundation.modules.direct_statement.direct_statement import DirectStatement
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.modules.direct_statement.direct_statement_repository import StatementRepository
from analogai.foundation.infrastructure.storage.in_memory.in_memory_storage_adapter_impl import InMemoryStorageAdapterImpl
from analogai.foundation.infrastructure.mappers.statement_storage_mapper import StatementStorageMapper
from analogai.foundation.common.utils.modifier_utils import ModifierMatcher


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 InMemoryStatementRepositoryImpl(StatementRepository):
    COLLECTION_NAME = "statements"

    def __init__(self):
        self.storage_adapter = InMemoryStorageAdapterImpl()
        self._modal_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: Relation, query_relation: Relation) -> bool:
        if not _use_vector_search() or not self._embedding_service:
            return candidate_relation == query_relation
        query_embedding = self._embedding_service.encode(query_relation.value)
        candidate_embedding = self._embedding_service.encode(candidate_relation.value)
        score = self._cosine_similarity(query_embedding, candidate_embedding)
        return score >= _similarity_cutoff()

    def create(self, statement: DirectStatement) -> DirectStatement:
        data = StatementStorageMapper.to_dict(statement)
        stored_data = self.storage_adapter.store(self.COLLECTION_NAME, data)
        if not stored_data:
            raise ValueError(f"Failed to store statement {statement.id}")
        return statement

    def update(self, statement: DirectStatement) -> DirectStatement:
        data = StatementStorageMapper.to_dict(statement)
        self.storage_adapter.store(self.COLLECTION_NAME, data)
        return statement

    def delete(self, statement: DirectStatement) -> None:
        if statement and statement.id:
            self.storage_adapter.delete(self.COLLECTION_NAME, statement.id)

    def find_by_id(self, statement_id: str) -> Optional[DirectStatement]:
        data = self.storage_adapter.retrieve_by_id(self.COLLECTION_NAME, statement_id)
        if not data:
            return None
        return StatementStorageMapper.from_dict(data)

    def find_by_subject(self, subject_notion: Notion) -> List[DirectStatement]:
        results = self.storage_adapter.retrieve_by_attributes(
            self.COLLECTION_NAME,
            {"subject_notion.id": subject_notion.id},
        )
        return [StatementStorageMapper.from_dict(d) for d in results]

    def find_by_subject_and_relation(
        self,
        subject_notion: Notion,
        relation: Relation,
        modifier: Optional[Modifier] = None,
    ) -> List[DirectStatement]:
        statements = self.find_by_subject(subject_notion)
        by_type = [p for p in statements if self._relation_matches(p.relation, relation)]
        if modifier is None:
            return by_type
        return [p for p in by_type if self._modal_matcher.context_matches(p.modifier, modifier)]

    def find_by_object(self, complement_notion: Notion) -> List[DirectStatement]:
        results = self.storage_adapter.retrieve_by_attributes(
            self.COLLECTION_NAME,
            {"complement_notion.id": complement_notion.id},
        )
        return [StatementStorageMapper.from_dict(d) for d in results]

    def find_by_object_and_relation(
        self,
        complement_notion: Notion,
        relation: Relation,
        modifier: Optional[Modifier] = None,
    ) -> List[DirectStatement]:
        statements = self.find_by_object(complement_notion)
        by_type = [p for p in statements if self._relation_matches(p.relation, relation)]
        if modifier is None:
            return by_type
        return [p for p in by_type if self._modal_matcher.context_matches(p.modifier, modifier)]

    def find_by_agent_id(
        self,
        agent_id: str,
        limit: int = 0,
        offset: int = 0,
    ) -> List[DirectStatement]:
        results = self.storage_adapter.retrieve_by_attributes(
            self.COLLECTION_NAME,
            {"agent_id": agent_id},
        )
        statements = [StatementStorageMapper.from_dict(d) for d in results]
        if limit <= 0:
            return statements[offset:]
        return statements[offset:offset + limit]

    def find_direct_statement(
        self,
        user_id: str,
        agent_id: str,
        subject_notion: Notion,
        relation: Relation,
        complement_notion: Notion,
        modifier: Optional[Modifier] = None,
    ) -> Optional[DirectStatement]:
        statements = self.find_by_subject_and_relation(
            subject_notion, relation, modifier
        )
        for stmt in statements:
            if (
                stmt.complement_notion.id == complement_notion.id
                and stmt.user_id == user_id
                and stmt.agent_id == agent_id
            ):
                return stmt
        return None

    def find_by_captions(
        self,
        agent_id: str,
        captions: List[str],
        max_depth: int = 1,
    ) -> List[DirectStatement]:
        if not captions:
            return []
        depth = max(1, int(max_depth))
        seed = {caption.strip().lower() for caption in captions if caption and caption.strip()}
        if not seed:
            return []
        all_statements = self.find_by_agent_id(agent_id)
        results: List[DirectStatement] = []
        seen_statement_ids = set()
        seen_captions = set(seed)
        frontier = set(seed)
        for _ in range(depth):
            if not frontier:
                break
            next_frontier = set()
            for stmt in all_statements:
                if stmt.id in seen_statement_ids:
                    continue
                subj = (stmt.subject_notion.caption or "").strip().lower()
                obj = (stmt.complement_notion.caption or "").strip().lower()
                if subj in frontier or obj in frontier:
                    results.append(stmt)
                    seen_statement_ids.add(stmt.id)
                    if subj:
                        next_frontier.add(subj)
                    if obj:
                        next_frontier.add(obj)
            next_frontier -= seen_captions
            seen_captions.update(next_frontier)
            frontier = next_frontier
        return results