# citation_handler.py

from langchain_core.documents import Document
from typing import Dict, List, Union, Any
import re
from .utilties.search_utilities import remove_think_tags
from .config import settings

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 and always keep URLS. Never make up sources. Example format: "According to the research [1], ..."
"""

        response = self.llm.invoke(prompt)

        return {"content": remove_think_tags(response.content), "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 settings.GENERAL.ENABLE_FACT_CHECKING:
            fact_check_response = remove_think_tags(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": remove_think_tags(response.content), "documents": documents}