from __future__ import annotations

from typing import Iterable

from analogai.deepthink.application.value_objects.user_agent import UserAgent
from analogai.foundation.modules.fact.services.fact_service import FactServicePort
from analogai.foundation.modules.fact.enums.fact_search_mode import FactSearchMode
from analogai.foundation.modules.fact.fact_source import FactSource
from analogai.deepthink.messaging.question_answering.services.assembly.fact_retrieval_command import (
    GraphRagSearchCommand,
)


class GraphRagBackgroundInfoBuilder:
    def __init__(self, fact_service: FactServicePort | None = None) -> None:
        self._fact_service = fact_service
        self._retrieval_command = GraphRagSearchCommand(self._fact_service)
        self._last_relaxed_used = False

    def build_background(
        self,
        user_agent: UserAgent,
        question_text: str,
        entities: Iterable[str] | None,
        notions: Iterable | None,
        search_mode: FactSearchMode,
        fact_lookup: dict[str, str] | None = None,
        source_lookup: dict[str, FactSource] | None = None,
        context_lookup: dict[str, list[str]] | None = None,
        parallel_workers: int | None = None,
        search_timeout_seconds: float | None = None,
    ) -> dict[str, str]:
        results = self._retrieval_command.execute(
            user_agent,
            question_text,
            entities,
            notions,
            search_mode,
            lookup=fact_lookup,
            source_lookup=source_lookup,
            context_lookup=context_lookup,
            parallel_workers=parallel_workers,
            search_timeout_seconds=search_timeout_seconds,
        )
        self._last_relaxed_used = self._retrieval_command.was_relaxed()
        return results

    def was_relaxed(self) -> bool:
        return bool(self._last_relaxed_used)

    def build_statement_facts(
        self,
        user_agent: UserAgent,
        entities: Iterable[str] | None,
        notions: Iterable | None,
        search_mode: FactSearchMode,
        fact_lookup: dict[str, str] | None = None,
        source_lookup: dict[str, FactSource] | None = None,
        context_lookup: dict[str, list[str]] | None = None,
        parallel_workers: int | None = None,
        search_timeout_seconds: float | None = None,
    ) -> dict[str, str]:
        return self.build_background(
            user_agent=user_agent,
            question_text="",
            entities=entities,
            notions=notions,
            search_mode=search_mode,
            fact_lookup=fact_lookup,
            source_lookup=source_lookup,
            context_lookup=context_lookup,
            parallel_workers=parallel_workers,
            search_timeout_seconds=search_timeout_seconds,
        )

    def build_direct_facts(
        self,
        user_agent: UserAgent,
        question_text: str,
        search_mode: FactSearchMode,
        fact_lookup: dict[str, str] | None = None,
        source_lookup: dict[str, FactSource] | None = None,
        context_lookup: dict[str, list[str]] | None = None,
    ) -> dict[str, str]:
        return self.build_background(
            user_agent=user_agent,
            question_text=question_text,
            entities=[],
            notions=None,
            search_mode=search_mode,
            fact_lookup=fact_lookup,
            source_lookup=source_lookup,
            context_lookup=context_lookup,
        )


# Backward compatibility
FactBackgroundInfoBuilder = GraphRagBackgroundInfoBuilder
