import asyncio
import inspect
import logging
import warnings
from typing import Any

import networkx as nx
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate
from tenacity import retry, stop_after_attempt, wait_exponential

from nxlu.config import Framework, Intent, NxluConfig
from nxlu.constants import SYSTEM_PROMPT, AnalysisPrompts
from nxlu.explanation.classify import IntentClassifier, intent2prompt
from nxlu.explanation.interrogate import GraphInterrogator
from nxlu.io import load_algorithm_docs
from nxlu.processing.analyze import get_algorithm_function
from nxlu.utils.control import init_llm_model
from nxlu.utils.misc import sanitize_input

warnings.filterwarnings("ignore")


logger = logging.getLogger("nxlu")


__all__ = ["UnifiedLLMChain", "GraphExplainer"]

MAX_LLM_THREADS = 5  # limiting to 5 concurrent calls for now


class UnifiedLLMChain:
    """A unified wrapper to handle both LangChain and LlamaIndex LLMs."""

    def __init__(self, config: NxluConfig):
        """Initialize the UnifiedLLMChain with the appropriate LLM and framework."""
        self.config = config
        framework = config.get_llm_framework()
        self.llm = init_llm_model(config)

        self.system_prompt = SYSTEM_PROMPT

        if framework == Framework.LANGCHAIN:
            self.chat_prompt = ChatPromptTemplate.from_messages(
                [("system", self.system_prompt), ("user", "{prompt}")]
            )
            self.chain = LLMChain(llm=self.llm, prompt=self.chat_prompt)

    def llamaindex_complete(self, prompt: str) -> str:
        """Structure the system and user prompts and call the LlamaIndex LLM."""
        full_prompt = f"{self.system_prompt}\n\nUser Prompt: {prompt}"
        return self.llm.complete(full_prompt)

    async def _call_llm_chain(self, prompt: str) -> str | None:
        """Call the LLM chain depending on the framework."""
        framework = self.config.get_llm_framework()
        if framework == Framework.LANGCHAIN:
            return await self.chain.arun(prompt=prompt)
        if framework == Framework.LLAMAINDEX:
            return self.llamaindex_complete(prompt)

    def run(self, prompt: str) -> str | None:
        """Run the chain synchronously depending on the framework."""
        framework = self.config.get_llm_framework()
        if framework == Framework.LANGCHAIN:
            return self.chain.run(prompt=prompt)
        if framework == Framework.LLAMAINDEX:
            return self.llamaindex_complete(prompt)


class GraphExplainer:
    """A class to generate explanations for graphs based on user queries using LLM and
    graph interrogation.

    Parameters
    ----------
    config : NxluConfig
        The configuration object used to initialize the LLM and GraphInterrogator.
    """

    def __init__(self, config: NxluConfig):
        """Initialize the GraphExplainer with LLM and GraphInterrogator."""
        self.config = config
        self.llm = init_llm_model(self.config)
        self.intent_classifier = None
        self.interrogator = None
        self._llm_semaphore = asyncio.Semaphore(MAX_LLM_THREADS)

        self.chain = UnifiedLLMChain(self.config)

    @retry(
        stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10)
    )
    async def _call_llm_chain(self, prompt: str) -> str | None:
        async with self._llm_semaphore:
            return await self.chain._call_llm_chain(prompt)

    async def explain_async(
        self,
        graph_data: nx.Graph,
        query: str | None = None,
    ) -> str:
        """Generate an explanation for the provided graph and query.

        Parameters
        ----------
        graph_data : nx.Graph
            The graph data to be analyzed.
        query : str
            The user's query to guide the explanation.

        Returns
        -------
        str
            The explanation generated by the LLM.
        """
        self.interrogator = GraphInterrogator(self.config)
        if query:
            query = sanitize_input(query)
            if not self.intent_classifier:
                self.intent_classifier = IntentClassifier()
            intent = self.intent_classifier.fit_transform(query)
            graph_info = await self.interrogator.reason_async(graph_data, query, intent)
            prompt = (
                f"Network Insights:"
                f"\nGraph Summary: {graph_info.get('Graph Summary')}\n"
                f"Descriptions of Applied Algorithms: "
                f"{graph_info.get('Descriptions of Applied Algorithms')}\n"
                f"Graph Analysis: {graph_info.get('Graph Analysis')}\n"
                f"Graph Relationships: {graph_info.get('Graph Relationships')}"
                f"\n\n**Query: {query}**\n\n"
                f"Please provide a response to the query based on the network "
                f"insights. Please avoid detailed technical descriptions unless "
                f"they help to shed light on your network-informed reasoning "
                f"mechanisms. Unless it pertains directly to the user's query, avoid "
                f"overinterpreting prominent nodes, as their hubness and "
                f"authoritativeness may merely serve to confound the analysis as a "
                f"whole and might ultimately be irrelevant to formulating a network "
                f"interpretation directed at the user's query."
            )
            logger.debug(f"Prompt:\n{prompt}")
        else:
            intent = [Intent.EXPLORATION]
            query = AnalysisPrompts.get_prompt(intent2prompt(intent))
            graph_info = await self.interrogator.reason_async(
                graph_data,
                query,
                intent,
            )
            prompt = (
                f"Network Insights:"
                f"\nGraph Summary: {graph_info.get('Graph Summary')}\n"
                f"Descriptions of Applied Algorithms: "
                f"{graph_info.get('Descriptions of Applied Algorithms')}\n"
                f"Graph Analysis: {graph_info.get('Graph Analysis')}\n"
                f"Graph Relationships: {graph_info.get('Graph Relationships')}\n\n"
                f"Please provide a comprehensive interpretation of this network."
            )
            logger.debug(f"Prompt:\n{prompt}")
        response = await self._call_llm_chain(prompt)
        return response.strip()

    def explain_algorithm(
        self, graph_data: nx.Graph, algorithm_name: str, result: Any
    ) -> str:
        """Generate a natural language explanation for a specific algorithm's results.

        Parameters
        ----------
        graph_data : networkx.Graph
            The graph on which the algorithm was run.
        algorithm_name : str
            The name of the algorithm.
        result : Any
            The result obtained from running the algorithm.

        Returns
        -------
        str
            The generated explanation.

        Raises
        ------
        Exception
            If explanation generation fails.
        """
        try:
            algorithm_docstring = inspect.getdoc(get_algorithm_function(algorithm_name))
            algorithm_description = load_algorithm_docs(algorithm_name)

            prompt = (
                f"Algorithm: {algorithm_name}\n\n"
                f"Algorithm Description:\n{algorithm_description}\n\n"
                f"Algorithm Docstring:\n{algorithm_docstring}\n\n"
                f"Result:\n{result}\n\n"
                f"Provide a clear and insightful explanation of the results of the "
                f"{algorithm_name} algorithm, considering its description and the "
                f"obtained result."
            )
            response = self.chain.run(prompt)
            return response.strip()
        except Exception:
            logger.exception(f"Error generating explanation for {algorithm_name}")
            return "Sorry, I couldn't generate an explanation at this time."

    def explain(self, graph_data: nx.Graph, query: str | None = None) -> str:
        """Generate an explanation by running the async method.

        Parameters
        ----------
        graph_data : nx.Graph
            The graph data to be analyzed.
        query : Optional[str]
            The user's query to guide the explanation.

        Returns
        -------
        str
            The explanation generated by the LLM.
        """
        return asyncio.run(self.explain_async(graph_data, query))
