from __future__ import annotations

from typing import Iterable, List, Tuple
from collections import deque
from concurrent.futures import ThreadPoolExecutor
from threading import Lock

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

from analogai.deepthink.application.value_objects.user_agent import UserAgent
from analogai.foundation.settings import config as app_config


class BeliefSearchAdapter:
    def __init__(
        self,
        belief_service,
        max_entity_depth: int | None = None,
        max_beliefs: int | None = None,
        parallel_workers: int | None = None,
        combined_query_enabled: bool | None = None,
        batch_query_enabled: bool | None = None,
    ) -> None:
        self._belief_service = belief_service
        resolved_entity_depth = (
            app_config.get(
                "qa.fact_building.max_entity_depth",
                app_config.get("qa.max_entity_depth", 3),
            )
            if max_entity_depth is None
            else max_entity_depth
        )
        resolved_max_beliefs = (
            app_config.get(
                "qa.fact_building.max_beliefs",
                app_config.get("qa.max_beliefs", 12),
            )
            if max_beliefs is None
            else max_beliefs
        )
        resolved_parallel_workers = (
            app_config.get("qa.belief_search_parallel_workers", 0)
            if parallel_workers is None
            else parallel_workers
        )
        resolved_combined_enabled = (
            app_config.get("qa.belief_search_combined_enabled", True)
            if combined_query_enabled is None
            else combined_query_enabled
        )
        resolved_batch_enabled = (
            app_config.get("qa.belief_search_batch_enabled", True)
            if batch_query_enabled is None
            else batch_query_enabled
        )
        self._max_entity_depth = max(0, int(resolved_entity_depth))
        self._max_beliefs = None if resolved_max_beliefs is None else max(0, int(resolved_max_beliefs))
        self._parallel_workers = max(0, int(resolved_parallel_workers))
        self._combined_query_enabled = bool(resolved_combined_enabled)
        self._batch_query_enabled = bool(resolved_batch_enabled)
        self._belief_cache: dict[str, Tuple[List, List]] = {}
        self._cache_lock = Lock()

    def _fetch_beliefs_for_entity(self, user_agent: UserAgent, entity: str) -> Tuple[List, List]:
        key = entity.strip().lower()
        if not key:
            return [], []
        with self._cache_lock:
            cached = self._belief_cache.get(key)
        if cached is not None:
            return cached
        notion_expression = NotionExpression(caption=entity)
        if self._combined_query_enabled and hasattr(self._belief_service, "search_by_notion_expression"):
            try:
                combined = (
                    self._belief_service.search_by_notion_expression(
                        user_agent.agent_id,
                        notion_expression,
                    )
                    or []
                )
                result = (combined, [])
                with self._cache_lock:
                    self._belief_cache[key] = result
                return result
            except TypeError:
                pass
        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 []
        )
        result = (subject_beliefs, complement_beliefs)
        with self._cache_lock:
            self._belief_cache[key] = result
        return result

    def search_beliefs(self, user_agent: UserAgent, entities: Iterable[str]) -> List:
        if self._belief_service is None:
            return []
        if self._max_beliefs is not None and self._max_beliefs <= 0:
            return []
        if self._batch_query_enabled and hasattr(self._belief_service, "search_by_notion_captions"):
            captions = [entity for entity in entities if entity and str(entity).strip()]
            if captions:
                try:
                    return self._belief_service.search_by_notion_captions(
                        user_agent.agent_id,
                        captions,
                        max_depth=self._max_entity_depth,
                    )
                except TypeError:
                    return self._belief_service.search_by_notion_captions(
                        user_agent.agent_id,
                        captions,
                    )
        results: List = []
        seen_belief_ids = set()
        entity_queue = deque((entity, 0) for entity in entities if entity)
        seen_entity_keys = {entity.lower() for entity, _ in entity_queue}
        executor = (
            ThreadPoolExecutor(max_workers=self._parallel_workers)
            if self._parallel_workers > 1
            else None
        )
        try:
            while entity_queue:
                if self._max_beliefs is not None and len(results) >= self._max_beliefs:
                    break
                entity, depth = entity_queue.popleft()
                if depth > self._max_entity_depth:
                    continue

                batch = [(entity, depth)]
                while entity_queue and entity_queue[0][1] == depth:
                    batch.append(entity_queue.popleft())

                if executor is not None and len(batch) > 1:
                    futures = [
                        executor.submit(self._fetch_beliefs_for_entity, user_agent, item[0])
                        for item in batch
                    ]
                    batch_results = [future.result() for future in futures]
                else:
                    batch_results = [
                        self._fetch_beliefs_for_entity(user_agent, item[0])
                        for item in batch
                    ]

                for (subject_beliefs, complement_beliefs), (_, batch_depth) in zip(batch_results, batch):
                    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 self._max_beliefs is not None and len(results) >= self._max_beliefs:
                            break

                        if batch_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, batch_depth + 1))

                    if self._max_beliefs is not None and len(results) >= self._max_beliefs:
                        break
        finally:
            if executor is not None:
                executor.shutdown(wait=True)

        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
