"""
Base class for all citation handlers.
"""

from abc import ABC, abstractmethod
from typing import Any, Callable, Dict, List, Optional, Union

from langchain_core.documents import Document
from loguru import logger

from ..utilities.type_utils import unwrap_setting
from ..utilities.json_utils import get_llm_response_text


class BaseCitationHandler(ABC):
    """Abstract base class for citation handlers."""

    def __init__(self, llm, settings_snapshot=None):
        self.llm = llm
        self.settings_snapshot = settings_snapshot or {}
        self._fact_checking_logged = False
        self.stream_callback: Optional[Callable[[str], None]] = None

    def set_stream_callback(self, callback: Callable[[str], None]):
        """Set a callback that receives each streamed LLM token."""
        self.stream_callback = callback

    def _invoke_with_streaming(self, prompt: str) -> str:
        """
        Invoke the LLM, streaming tokens through the callback if set.

        Falls back to a single ``invoke()`` call when no callback is
        registered or when the LLM does not support ``.stream()``.

        Returns:
            The complete response text.
        """
        if self.stream_callback and hasattr(self.llm, "stream"):
            chunks = []
            try:
                for chunk in self.llm.stream(prompt):
                    text = (
                        chunk
                        if isinstance(chunk, str)
                        else getattr(chunk, "content", str(chunk))
                    )
                    if text:
                        chunks.append(text)
                        try:
                            self.stream_callback(text)
                        except Exception:
                            logger.debug(
                                "stream_callback failed"
                            )  # Non-critical: don't break synthesis
                # Normalize the joined chunks exactly like the invoke() path
                # below: .stream() bypasses ProcessingLLMWrapper.invoke (which
                # is the only place <think> blocks are stripped), so without
                # this a reasoning model's <think>…</think> would leak into the
                # persisted answer. The live token stream is a separate concern.
                return get_llm_response_text("".join(chunks))
            except Exception:
                # If any chunks already crossed the wire to the client,
                # restarting via .invoke() would (a) double-bill the LLM
                # and (b) cause the frontend's accumulated streamed text
                # to diverge from the new full response — the chat bubble
                # then shows partial chunks while the DB row carries the
                # invoke()-result. Only fall back when nothing was emitted.
                if chunks:
                    logger.warning(
                        "Stream errored after {} chunks; returning partial "
                        "content (no invoke() fallback to avoid double-bill "
                        "and UI/DB divergence)",
                        len(chunks),
                    )
                    return get_llm_response_text("".join(chunks))
                logger.debug(
                    "Stream failed before any chunk; falling back to invoke()"
                )

        # No callback (non-chat research) or stream unavailable: delegate to
        # the same normalization _invoke_text uses, so <think> blocks
        # are stripped and str/object responses are handled uniformly.
        return get_llm_response_text(self.llm.invoke(prompt))

    def _invoke_text(self, prompt: str) -> str:
        """Invoke the LLM and return normalized text.

        Handles both message objects (``.content``) and raw string responses,
        and strips ``<think>`` reasoning blocks via ``get_llm_response_text``.
        """
        return get_llm_response_text(self.llm.invoke(prompt))

    def get_setting(self, key: str, default=None):
        """Get a setting value from the snapshot."""
        if key in self.settings_snapshot:
            return unwrap_setting(self.settings_snapshot[key])
        return default

    def is_fact_checking_enabled(self) -> bool:
        """Check if fact-checking is enabled and log the state once."""
        enabled = self.get_setting("general.enable_fact_checking", False)
        if not self._fact_checking_logged:
            handler_name = type(self).__name__
            if enabled:
                logger.info(
                    f"[{handler_name}] Fact-checking is ENABLED — "
                    f"extra LLM call per synthesis"
                )
            else:
                logger.info(f"[{handler_name}] Fact-checking is DISABLED")
            self._fact_checking_logged = True
        return bool(enabled)

    def _get_output_instruction_prefix(self) -> str:
        """
        Get formatted output instructions from settings if present.

        This allows users to customize output language, tone, style, and formatting
        for research answers and reports. Instructions are prepended to prompts
        sent to the LLM.

        Returns:
            str: Formatted instruction prefix if custom instructions are set,
                 empty string otherwise.

        Examples:
            - "Respond in Spanish with formal academic tone"
            - "Use simple language suitable for beginners"
            - "Be concise with bullet points"
        """
        output_instructions = self.get_setting(
            "general.output_instructions", ""
        ).strip()

        if output_instructions:
            return f"User-Specified Output Style: {output_instructions}\n\n"
        return ""

    def _create_documents(
        self, search_results: Union[str, List[Dict]], nr_of_links: int = 0
    ) -> List[Document]:
        """
        Convert search results to LangChain documents format and add index
        to original search results.
        """
        documents: List[Document] = []
        if isinstance(search_results, str):
            return documents

        for i, result in enumerate(search_results):
            if isinstance(result, dict):
                # Add index to the original search result dictionary if it doesn't exist
                # This preserves indices that were already set (e.g., for topic organization)
                if "index" not in result:
                    result["index"] = str(i + nr_of_links + 1)

                content = result.get("full_content", result.get("snippet", ""))
                # Use the index from the result if it exists, otherwise calculate it
                doc_index = int(result.get("index", i + nr_of_links + 1))
                documents.append(
                    Document(
                        page_content=content,
                        metadata={
                            "source": result.get("link", f"source_{i + 1}"),
                            "title": result.get("title", f"Source {i + 1}"),
                            "index": doc_index,
                        },
                    )
                )
        return documents

    def _format_sources(self, documents: List[Document]) -> str:
        """Format sources with numbers for citation."""
        sources = []
        for doc in documents:
            source_id = doc.metadata["index"]
            sources.append(f"[{source_id}] {doc.page_content}")
        return "\n\n".join(sources)

    def _no_sources_response(self, question: str) -> Dict[str, Any]:
        """
        Explicit no-sources result returned instead of invoking the LLM.

        Prompting the LLM to "answer with citations [1], [2]…" while the
        sources section is empty makes it fall back on its training data
        and fabricate references. Handlers call this to refuse synthesis
        instead.
        """
        logger.warning(
            f"[{type(self).__name__}] No sources available for synthesis of "
            f"'{question[:100]}' — skipping LLM call to avoid fabricated "
            f"citations"
        )
        content = (
            "No sources were found for this question. The selected search "
            "engines or document collections returned no results; this can "
            "also happen when a search fails with an error (check the "
            "research logs). No answer was generated because, without "
            "sources, it would have to rely on the language model's "
            "built-in knowledge and could contain fabricated citations."
        )
        return {"content": content, "documents": []}

    @abstractmethod
    def analyze_initial(
        self, query: str, search_results: Union[str, List[Dict]]
    ) -> Dict[str, Any]:
        """Process initial analysis with citations."""
        pass

    @abstractmethod
    def analyze_followup(
        self,
        question: str,
        search_results: Union[str, List[Dict]],
        previous_knowledge: str,
        nr_of_links: int,
    ) -> Dict[str, Any]:
        """Process follow-up analysis with citations."""
        pass
