from typing import Optional, List
import math

from analogai.foundation import config
from analogai.foundation.common.logging.app_logger import app_logger
from analogai.foundation.modules.notion.notion_repository import NotionRepository
from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.common.dtos.notion_expression import NotionExpression
from analogai.foundation.infrastructure.storage.in_memory.in_memory_storage_adapter_impl import InMemoryStorageAdapterImpl
from analogai.foundation.infrastructure.mappers.notion_storage_mapper import NotionStorageMapper


class InMemoryNotionRepositoryImpl(NotionRepository):
    COLLECTION_NAME = "notions"

    def __init__(self, storage_adapter=None):
        if storage_adapter is None:
            self.storage_adapter = InMemoryStorageAdapterImpl()
        else:
            self.storage_adapter = storage_adapter
        self._embedding_service = None
        if config.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 create(self, notion: Notion) -> Notion:
        data = NotionStorageMapper.to_dict(notion)
        if config.USE_VECTOR_SEARCH and self._embedding_service:
            data["embedding"] = self._embedding_service.encode(notion.caption)
        stored_data = self.storage_adapter.store(self.COLLECTION_NAME, data)
        if not stored_data:
            raise ValueError(f"Failed to store notion {notion.id}")
        return notion

    def find(
        self,
        agent_id: str,
        expression: NotionExpression,
        context_ids: Optional[List[str]] = None,
        similarity_threshold: Optional[float] = None,
    ) -> Optional[Notion]:
        if not config.USE_VECTOR_SEARCH or not self._embedding_service:
            query = {"caption": expression.caption, "agent_id": agent_id}
            all_notions = self.storage_adapter.retrieve_by_attributes(self.COLLECTION_NAME, query)

            if not all_notions:
                return None

            notion_data = all_notions[0]
            notion = NotionStorageMapper.from_dict(notion_data)
            return notion

        candidates = self.storage_adapter.retrieve_by_attributes(
            self.COLLECTION_NAME,
            {"agent_id": agent_id},
        )
        if not candidates:
            return None

        threshold = (
            similarity_threshold
            if similarity_threshold is not None
            else config.SIMILARITY_CUTOFF
        )
        query_embedding = self._embedding_service.encode(expression.caption)
        best_match = None
        best_score = 0.0
        for notion_data in candidates:
            embedding = notion_data.get("embedding")
            if embedding is None:
                embedding = self._embedding_service.encode(notion_data.get("caption", ""))
            score = self._cosine_similarity(query_embedding, embedding)
            app_logger.debug(
                "Vector similarity notion match: query='%s' candidate='%s' score=%.4f threshold=%.4f",
                expression.caption,
                notion_data.get("caption", ""),
                score,
                threshold,
            )
            if score >= threshold and score > best_score:
                best_score = score
                best_match = notion_data

        if not best_match:
            return None

        app_logger.debug(
            "Vector notion match selected: query='%s' candidate='%s' score=%.4f",
            expression.caption,
            best_match.get("caption", ""),
            best_score,
        )
        return NotionStorageMapper.from_dict(best_match)

    def find_by_id(self, notion_id: str) -> Optional[Notion]:
        result_data = self.storage_adapter.retrieve_by_id(self.COLLECTION_NAME, notion_id)
        if result_data:
            notion = NotionStorageMapper.from_dict(result_data)
            return notion
        return None

    def find_all_by_agent(self, agent_id: str) -> List[Notion]:
        results = self.storage_adapter.retrieve_by_attributes(
            self.COLLECTION_NAME,
            {"agent_id": agent_id},
        )
        return [NotionStorageMapper.from_dict(d) for d in results]

