from __future__ import annotations

from typing import Iterable, List, Optional

from analogai.foundation.common.enums.criterion import Criterion
from analogai.foundation.common.utils.expression_factory import ExpressionFactory
from analogai.foundation.common.dtos.notion_expression import NotionExpression

from analogai.deepthink.application.value_objects.user_agent import UserAgent
from analogai.deepthink.features.question_answering.question_entity_extractor import QuestionEntityExtractor
from analogai.deepthink.features.question_answering.question_llm_caller import QuestionLlmCaller
from analogai.deepthink.features.question_answering.question_prompt_builder import (
    QuestionPromptBuilder,
    QuestionPromptContext,
    QuestionSynthesisContext,
)
from analogai.deepthink.infrastructure.gateways.llm.llm_gateway_provider_factory import (
    LlmGatewayProviderFactory,
)

try:
    from analogai.foundation.common.utils.expressions_serializer import serialize_belief_expression
except Exception:  # pragma: no cover - optional dependency
    serialize_belief_expression = None


class QuestionAnsweringService:
    def __init__(
        self,
        belief_query_service,
        notion_service,
        llm_gateway,
        synthesizer_gateway=None,
        entity_extractor: Optional[QuestionEntityExtractor] = None,
        prompt_builder: Optional[QuestionPromptBuilder] = None,
        llm_caller: Optional[QuestionLlmCaller] = None,
        synthesizer_caller: Optional[QuestionLlmCaller] = None,
        max_entity_depth: int = 5,
        similarity_cutoff: Optional[float] = None,
        use_vector_search: Optional[bool] = None,
    ):
        self._belief_query_service = belief_query_service
        self._notion_service = notion_service
        self._entity_extractor = entity_extractor or QuestionEntityExtractor()
        self._prompt_builder = prompt_builder or QuestionPromptBuilder()
        self._llm_caller = llm_caller or QuestionLlmCaller(llm_gateway)
        self._synthesizer_gateway = synthesizer_gateway
        self._synthesizer_caller = synthesizer_caller
        self._expression_factory = ExpressionFactory()
        self._max_entity_depth = max(0, int(max_entity_depth))
        self._similarity_cutoff = None
        if similarity_cutoff is not None:
            try:
                cutoff = float(similarity_cutoff)
                cutoff = max(0.0, min(1.0, cutoff))
                self._similarity_cutoff = cutoff
            except (TypeError, ValueError):
                self._similarity_cutoff = None
        self._use_vector_search = None
        if use_vector_search is not None:
            self._use_vector_search = bool(use_vector_search)

    def search(
        self,
        question_text: str,
        user_agent: UserAgent,
    ) -> Optional[str]:
        if not question_text or not question_text.strip():
            return None

        prompt_context, belief_count, sentence_count = self._build_prompt_context_with_counts(
            question_text=question_text,
            user_agent=user_agent,
        )
        if prompt_context is None:
            return None
        if belief_count + sentence_count == 0:
            return None
        user_message = self._prompt_builder.build_user_message(prompt_context)
        raw_answer = self._llm_caller.answer(user_message)

        synthesized = self._synthesize_answer(
            question_text=question_text.strip(),
            belief_context=prompt_context.belief_context,
            raw_answer=raw_answer,
        )
        return synthesized

    def generate_context(
        self,
        question_text: str,
        user_agent: UserAgent,
    ) -> str:
        if not question_text or not question_text.strip():
            return ""

        prompt_context, _, _ = self._build_prompt_context_with_counts(
            question_text=question_text,
            user_agent=user_agent,
        )
        if prompt_context is None:
            return ""
        return prompt_context.belief_context

    def get_context(
        self,
        question_text: str,
        user_agent: UserAgent,
    ) -> str:
        return self.generate_context(
            question_text=question_text,
            user_agent=user_agent,
        )

    def _synthesize_answer(
        self,
        question_text: str,
        belief_context: str,
        raw_answer: str,
    ) -> str:
        facts = _build_synth_facts(belief_context, raw_answer)
        if not facts:
            return ""
        synthesis_context = QuestionSynthesisContext(
            question=question_text,
            tasks="",
            facts=facts,
        )
        prompt = self._prompt_builder.build_synthesis_prompt(synthesis_context)
        caller = self._get_synthesizer_caller()
        return caller.answer(prompt)

    def _build_prompt_context_with_counts(
        self,
        question_text: str,
        user_agent: UserAgent,
    ) -> tuple[Optional[QuestionPromptContext], int, int]:
        entities = self._entity_extractor.extract(question_text)
        if not entities:
            return None, 0, 0
        beliefs = self._fetch_beliefs(user_agent, entities)
        notions = self._fetch_notions(user_agent, entities)
        notions = self._extend_notions_from_beliefs(user_agent, notions, beliefs)
        if not beliefs and not notions:
            return None, 0, 0
        belief_context = self._build_context_text(beliefs, notions)
        if not belief_context.strip():
            return None, 0, 0

        belief_count = self._count_serialized_beliefs(beliefs)
        sentence_count = self._count_context_lines(belief_context)

        return (
            QuestionPromptContext(
                question=question_text.strip(),
                entities=entities,
                belief_context=belief_context,
            ),
            belief_count,
            sentence_count,
        )

    def _count_serialized_beliefs(self, beliefs: Iterable) -> int:
        count = 0
        for belief in beliefs or []:
            if self._format_belief(belief):
                count += 1
        return count

    @staticmethod
    def _count_context_lines(text: str) -> int:
        return sum(1 for line in text.splitlines() if line.strip())
    def _get_synthesizer_caller(self) -> QuestionLlmCaller:
        if self._synthesizer_caller is not None:
            return self._synthesizer_caller
        if self._synthesizer_gateway is None:
            self._synthesizer_gateway = LlmGatewayProviderFactory.create_gateway(
                system_prompt="",
                max_tokens=200,
            )
        self._synthesizer_caller = QuestionLlmCaller(self._synthesizer_gateway)
        return self._synthesizer_caller

    def _fetch_beliefs(self, user_agent: UserAgent, entities: Iterable[str]) -> List:
        results: List = []
        seen = set()
        queue = [(entity, 0) for entity in entities if entity]
        seen_entities = {entity.lower() for entity, _ in queue}

        while queue:
            entity, depth = queue.pop(0)
            if depth > self._max_entity_depth:
                continue
            builder = self._belief_query_service.search(
                user_agent.user_id,
                user_agent.agent_id,
            ).where(
                Criterion.SubjectCaption, entity
            ).orWhere(
                Criterion.ComplementCaption, entity
            )
            self._apply_vector_search(builder)
            self._apply_similarity_cutoff(builder)
            try:
                beliefs = builder.list()
            except Exception:
                beliefs = []
            for belief in beliefs:
                key = getattr(belief, "id", None) or str(belief)
                if key in seen:
                    continue
                seen.add(key)
                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_entities:
                            continue
                        seen_entities.add(related_key)
                        queue.append((related, depth + 1))

        return results

    def _apply_similarity_cutoff(self, builder) -> None:
        if self._similarity_cutoff is None:
            return
        for method_name in (
            "similarity_cutoff",
            "with_similarity_cutoff",
            "set_similarity_cutoff",
            "set_similarity_threshold",
            "similarity_threshold",
            "set_similarity_score_threshold",
            "set_min_similarity",
            "min_similarity",
        ):
            method = getattr(builder, method_name, None)
            if callable(method):
                method(self._similarity_cutoff)
                return

    def _apply_vector_search(self, builder) -> None:
        if self._use_vector_search is None:
            return
        for method_name in (
            "use_vector_search",
            "with_vector_search",
            "set_use_vector_search",
            "set_vector_search",
            "set_vector_search_enabled",
            "enable_vector_search",
            "vector_search",
            "use_vector",
        ):
            method = getattr(builder, method_name, None)
            if callable(method):
                try:
                    method(self._use_vector_search)
                except TypeError:
                    method()
                return

    def _format_beliefs(self, beliefs: Iterable) -> str:
        formatted = []
        for belief in beliefs:
            formatted.append(self._format_belief(belief))
        return "\n".join([line for line in formatted if line])

    def _format_notions(self, notions: Iterable) -> str:
        formatted: List[str] = []
        for notion in notions:
            try:
                formatted.append(str(notion))
            except Exception:
                continue
        return "\n".join([line for line in formatted if line])

    def _build_context_text(self, beliefs: Iterable, notions: Iterable) -> str:
        belief_text = self._format_beliefs(beliefs) if beliefs else ""
        notion_text = self._format_notions(notions) if notions else ""
        if belief_text and notion_text:
            return f"{belief_text}\n{notion_text}"
        return belief_text or notion_text

    def _format_belief(self, belief) -> str:
        try:
            return str(belief)
        except Exception:
            pass
        try:
            expression = self._expression_factory.propagate_belief_expression(belief)
            return _format_expression(expression)
        except Exception:
            return str(belief)

    def _fetch_notions(self, user_agent: UserAgent, entities: Iterable[str]) -> List:
        if self._notion_service is None:
            return []
        results: List = []
        seen = set()
        for entity in entities:
            if not entity:
                continue
            key = entity.strip().lower()
            if key in seen:
                continue
            seen.add(key)
            try:
                notion = self._notion_service.find(
                    user_agent.agent_id,
                    NotionExpression(caption=entity),
                )
            except Exception:
                notion = None
            if notion is not None:
                results.append(notion)
        return results

    def _extend_notions_from_beliefs(self, user_agent: UserAgent, notions: List, beliefs: Iterable) -> List:
        if self._notion_service is None:
            return notions
        results = list(notions) if notions else []
        seen = {getattr(notion, "id", None) or getattr(notion, "caption", "") for notion in results}
        for belief in beliefs or []:
            for notion in (getattr(belief, "subject_notion", None), getattr(belief, "complement_notion", None)):
                if notion is None:
                    continue
                notion_id = getattr(notion, "id", None)
                notion_caption = getattr(notion, "caption", None)
                key = notion_id or notion_caption or str(notion)
                if key in seen:
                    continue
                if notion_caption:
                    try:
                        fetched = self._notion_service.find(
                            user_agent.agent_id,
                            NotionExpression(caption=notion_caption),
                        )
                    except Exception:
                        fetched = None
                    if fetched is not None:
                        results.append(fetched)
                        seen.add(getattr(fetched, "id", None) or getattr(fetched, "caption", "") or key)
                        continue
                results.append(notion)
                seen.add(key)
        return results


def _format_expression(expression) -> str:
    subject_expr = getattr(expression, "subject_notion_expression", None) or getattr(expression, "subject", None)
    complement_expr = getattr(expression, "complement_notion_expression", None) or getattr(expression, "complement", None)
    relation = getattr(expression, "relation", None)
    probability = getattr(expression, "probability", None)
    modifier = getattr(expression, "modifier", None)

    subject = _caption(subject_expr)
    complement = _caption(complement_expr)
    relation_value = getattr(relation, "value", None) or str(relation) if relation is not None else ""

    parts = [p for p in [subject, relation_value, complement] if p]
    base = " ".join(parts) if parts else str(expression)

    modifier_bits = []
    if modifier is not None:
        location = getattr(modifier, "location", None)
        time_value = getattr(modifier, "time", None)
        quantity = getattr(modifier, "quantity", None)
        deontic = getattr(modifier, "deontic_modality", None)
        if location:
            modifier_bits.append(f"location={location}")
        if time_value:
            modifier_bits.append(f"time={time_value}")
        if quantity is not None:
            modifier_bits.append(f"quantity={quantity}")
        if deontic:
            modifier_bits.append(f"deontic={deontic}")

    if probability is not None:
        modifier_bits.append(f"probability={probability}")

    if modifier_bits:
        return f"{base} ({', '.join(modifier_bits)})"
    return base


def _caption(value) -> str:
    if value is None:
        return ""
    caption = getattr(value, "caption", None)
    attributes = _extract_notion_attributes(value)
    if caption:
        if attributes:
            return f"{caption} ({_format_attributes(attributes)})"
        return str(caption)
    if attributes:
        return _format_attributes(attributes)
    return str(value)


def _extract_belief_entities(belief) -> List[str]:
    entities = []
    subject_notion = getattr(belief, "subject_notion", None)
    complement_notion = getattr(belief, "complement_notion", None)
    if subject_notion is not None:
        caption = getattr(subject_notion, "caption", None)
        if caption:
            entities.append(str(caption))
    if complement_notion is not None:
        caption = getattr(complement_notion, "caption", None)
        if caption:
            entities.append(str(caption))
    return entities


def _extract_notion_attributes(notion) -> list:
    if notion is None:
        return []
    attributes = getattr(notion, "attributes", None)
    if attributes is None:
        return []
    if isinstance(attributes, list):
        return attributes
    return list(attributes) if attributes else []


def _format_notion_attributes(notion) -> str:
    caption = getattr(notion, "caption", None)
    attributes = _extract_notion_attributes(notion)
    if not attributes:
        return ""
    attrs_text = _format_attributes(attributes)
    if caption:
        return f"{caption} ({attrs_text})"
    return attrs_text


def _format_attributes(attributes: list) -> str:
    parts: List[str] = []
    for attr in attributes:
        tag = getattr(attr, "tag", None) or (attr.get("tag") if isinstance(attr, dict) else None)
        value = getattr(attr, "value", None) if not isinstance(attr, dict) else attr.get("value")
        unit = getattr(attr, "unit", None) if not isinstance(attr, dict) else attr.get("unit")

        if value is None and hasattr(attr, "values"):
            values = getattr(attr, "values", None)
            if values:
                value = values

        if value is None:
            if tag:
                parts.append(str(tag))
            continue

        if isinstance(value, list):
            value_str = ", ".join(str(v) for v in value)
        else:
            value_str = str(value)

        unit_str = f" {unit}" if unit else ""
        if tag:
            parts.append(f"{tag}={value_str}{unit_str}")
        else:
            parts.append(f"{value_str}{unit_str}")
    return ", ".join(parts)


def _build_synth_facts(belief_context: str, raw_answer: str) -> str:
    lines: List[str] = []
    if belief_context:
        for line in belief_context.splitlines():
            trimmed = line.strip()
            if trimmed:
                lines.append(f"- [belief] {trimmed}")
    if raw_answer and raw_answer.strip():
        lines.append(f"- [qa] {raw_answer.strip()}")
    return "\n".join(lines)
