# citation_handler.py

from typing import Any, Dict, List, Union

from langchain_core.documents import Document

from .utilities.db_utils import get_db_setting


class CitationHandler:
    def __init__(self, llm):
        self.llm = llm

    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)

    def analyze_initial(
        self, query: str, search_results: Union[str, List[Dict]]
    ) -> Dict[str, Any]:

        documents = self._create_documents(search_results)
        formatted_sources = self._format_sources(documents)
        prompt = f"""Analyze the following information concerning the question and include citations using numbers in square brackets [1], [2], etc. When citing, use the source number provided at the start of each source.

Question: {query}

Sources:
{formatted_sources}

Provide a detailed analysis with citations. Do not create the bibliography, it will be provided automatically.  Never make up sources. Never write or create urls. Only write text relevant to the question. Example format: "According to the research [1], ..."
"""

        response = self.llm.invoke(prompt)
        if not isinstance(response, str):
            response = response.content
        return {"content": response, "documents": documents}

    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."""
        documents = self._create_documents(search_results, nr_of_links=nr_of_links)
        formatted_sources = self._format_sources(documents)
        # Add fact-checking step
        fact_check_prompt = f"""Analyze these sources for factual consistency:
1. Cross-reference major claims between sources
2. Identify and flag any contradictions
3. Verify basic facts (dates, company names, ownership)
4. Note when sources disagree

Previous Knowledge:
{previous_knowledge}

New Sources:
{formatted_sources}

        Return any inconsistencies or conflicts found."""
        if get_db_setting("general.enable_fact_checking", True):
            fact_check_response = self.llm.invoke(fact_check_prompt).content

        else:
            fact_check_response = ""

        prompt = f"""Using the previous knowledge and new sources, answer the question. Include citations using numbers in square brackets [1], [2], etc. When citing, use the source number provided at the start of each source. Reflect information from sources critically.

Previous Knowledge:
{previous_knowledge}

Question: {question}

New Sources:
{formatted_sources}
Reflect information from sources critically based on: {fact_check_response}. Never invent sources.
Provide a detailed answer with citations.  Example format: "According to [1], ..." """

        response = self.llm.invoke(prompt)

        return {"content": response.content, "documents": documents}
