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

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

from langchain_core.documents import Document


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 {}

    def get_setting(self, key: str, default=None):
        """Get a setting value from the snapshot."""
        if key in self.settings_snapshot:
            value = self.settings_snapshot[key]
            # Extract value from dict structure if needed
            if isinstance(value, dict) and "value" in value:
                return value["value"]
            return value
        return default

    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 = []
        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
                result["index"] = str(i + nr_of_links + 1)

                content = result.get("full_content", result.get("snippet", ""))
                documents.append(
                    Document(
                        page_content=content,
                        metadata={
                            "source": result.get("link", f"source_{i + 1}"),
                            "title": result.get("title", f"Source {i + 1}"),
                            "index": i + nr_of_links + 1,
                        },
                    )
                )
        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)

    @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
