from __future__ import annotations

from typing import Iterable, List

from analogai.foundation.common.dtos.notion_expression import NotionExpression  # type: ignore[import-not-found]

from analogai.deepthink.application.value_objects.user_agent import UserAgent


class BeliefSearchAdapter:
    def __init__(
        self,
        belief_service,
        max_entity_depth: int = 5,
    ) -> None:
        self._belief_service = belief_service
        self._max_entity_depth = max(0, int(max_entity_depth))

    def search_beliefs(self, user_agent: UserAgent, entities: Iterable[str]) -> List:
        if self._belief_service is None:
            return []
        results: List = []
        seen_belief_ids = set()
        entity_queue = [(entity, 0) for entity in entities if entity]
        seen_entity_keys = {entity.lower() for entity, _ in entity_queue}

        while entity_queue:
            entity, depth = entity_queue.pop(0)
            if depth > self._max_entity_depth:
                continue
            notion_expression = NotionExpression(caption=entity)
            subject_beliefs = (
                self._belief_service.search_by_subject_notion_expression(
                    user_agent.agent_id,
                    notion_expression,
                )
                or []
            )
            complement_beliefs = (
                self._belief_service.search_by_complement_notion_expression(
                    user_agent.agent_id,
                    notion_expression,
                )
                or []
            )
            for belief in [*subject_beliefs, *complement_beliefs]:
                belief_id = belief.id
                if belief_id in seen_belief_ids:
                    continue
                seen_belief_ids.add(belief_id)
                results.append(belief)

                if depth < self._max_entity_depth:
                    for related in _extract_belief_entities(belief):
                        related_key = related.lower()
                        if related_key in seen_entity_keys:
                            continue
                        seen_entity_keys.add(related_key)
                        entity_queue.append((related, depth + 1))

        return results


def _extract_belief_entities(belief) -> List[str]:
    entities: List[str] = []
    subject_notion, complement_notion = _belief_notions(belief)
    subject_caption = _notion_caption(subject_notion)
    complement_caption = _notion_caption(complement_notion)
    if subject_caption:
        entities.append(str(subject_caption))
    if complement_caption:
        entities.append(str(complement_caption))
    return entities


def _belief_notions(belief) -> tuple[object | None, object | None]:
    if belief is None:
        return (None, None)
    return (belief.subject_notion, belief.complement_notion)


def _notion_caption(notion) -> str | None:
    if notion is None:
        return None
    return notion.caption
