from langchain.tools import tool
from langchain_community.tools import DuckDuckGoSearchResults
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from loguru import logger
from opentelemetry.trace import Status, StatusCode, get_tracer_provider
from rag import get_vector_store

tool_model: ChatOpenAI = None


def initialize_tool_llm(model):
    global tool_model
    tool_model = ChatOpenAI(model=model, temperature=0.7)
    return tool_model


@tool
def create_rag_response(user_query: str) -> str:
    """
    USE THIS TOOL FIRST for any user query. Fetches relevant information for the user's query from a pre-loaded vector store containing web content.

    This tool searches through embedded documents that were loaded from a web URL. Use this tool to retrieve context before answering any question.

    Args:
        user_query (str): The user's query or question.

    Returns:
        str: A concatenated string of relevant document contents retrieved from the vector store. Use this information to answer the user's question.
    """
    vector_store = get_vector_store()
    if vector_store is None:
        return "Error: Vector store is not initialized. Please ensure the web URL has been loaded."

    tracer = get_tracer_provider().get_tracer(__name__)

    with tracer.start_as_current_span("rag_retrieval", openinference_span_kind="retriever") as span:
        span.set_attribute("input.value", user_query)

        retrieved_docs = vector_store.similarity_search(user_query)
        logger.info(f"RAG search for '{user_query}' retrieved {len(retrieved_docs)} documents")

        if not retrieved_docs:
            span.set_attribute("retrieval.documents", 0)
            span.set_status(Status(StatusCode.OK))
            return "No relevant information found for the given query."

        # Set document attributes for each retrieved document
        for i, doc in enumerate(retrieved_docs):
            doc_id = str(i)
            doc_content = doc.page_content[:1200].replace("\n", " ")

            span.set_attribute(f"retrieval.documents.{i}.document.id", str(doc_id))
            span.set_attribute(f"retrieval.documents.{i}.document.content", doc_content)

        span.set_attribute("retrieval.documents", len(retrieved_docs))
        span.set_status(Status(StatusCode.OK))

        result = "\n\n".join(doc.page_content for doc in retrieved_docs)
        span.set_attribute("output.value", result)

        return result


@tool
def analyze_rag_response(query, rag_response: str) -> str:
    """
    This tool is designed to analyze the RAG (Retrieval-Augmented Generation) response generated for a user's query.

    When a user requests an analysis of the RAG response, this tool evaluates the response to ensure it matches the user's query.
    It provides insights into the response's relevance, clarity, and overall quality.

    Args:
        query (str): The original query provided by the user.
        rag_response (str): The response generated by the `create_rag_response` tool.

    Returns:
        str: A JSON analysis with:
        - **Original Response**: Full text of the RAG response.
        - **Key Points**: Main ideas in bullet points.
        - **Clarity**: A 1-10 clarity rating with reasoning.
        - **Relevance**: A 1-10 relevance rating.
        - **Suggestions**: Recommendations for improvement.

    Use Case:
    - Invoke when the user requests an analysis of the RAG response.
    - Compares the response to the query, highlights key points, and provides actionable feedback.
    """

    prompt = f"""
        You are a highly capable analysis tool specializing in evaluating responses generated by a Retrieval-Augmented Generation (RAG) system.
        Your task is to analyze the given response and provide the following details:
        1. **Original Response**: Include the full text of the response being analyzed.
        2. **Key Points**: Extract the main ideas or key points presented in the response. Provide these in a concise bullet-point format.
        3. **Clarity**: Assess the clarity of the response on a scale of 1 to 10, where 10 indicates perfect clarity. Briefly explain your reasoning.
        4. **Relevance**: Determine whether the response directly answers the query it was generated for. Rate relevance on a scale of 1 to 10.
        5. **Suggestions for Improvement**: If applicable, provide suggestions on how the response could be improved for better quality, clarity, or relevance.

        ### Input:
        Query: {query}
        RAG Response: {rag_response}

        ### Output:
        Provide your analysis in the following JSON format:
        {{
            "original_response": "{rag_response}",
            "key_points": ["Key point 1", "Key point 2", "Key point 3"],
            "clarity": "Clarity rating (1-10) - Explanation",
            "relevance": "Relevance rating (1-10)",
            "suggestions": ["Suggestion 1", "Suggestion 2"]
        }}
        """
    response = tool_model.invoke([HumanMessage(content=prompt)])
    return response.content


@tool
def web_search(user_query: str) -> str:
    """
    Performs a web search using DuckDuckGo to fetch the top 5 relevant results for the user's query.

    This tool is used when no relevant information is found in the RAG response. It retrieves
    the top 5 search results from DuckDuckGo, including titles, snippets, and URLs.

    Args:
        user_query (str): The user's query.

    Returns:
        str: A summary of the top 5 search results or a message indicating that no information was found.

    Example:
        Input:
            user_query = "Latest advancements in quantum computing."
        Output:
            "Top Search Results:
             1. Title: Quantum Computing Breakthroughs
                Snippet: Recent advancements in quantum computing include...
                URL: https://example.com/article1
             2. Title: The Future of Quantum Tech
                Snippet: Quantum technologies are set to revolutionize...
                URL: https://example.com/article2
             ...
            "
    """
    try:
        duck_duck_go = DuckDuckGoSearchResults(max_results=5, output_format="list")
        search_results = duck_duck_go.invoke(user_query)

        if not search_results:
            logger.info("Websearch completed, no documents retrieved")
            return "No relevant web search results were found for the given query."

        formatted_results = [
            f"{i}. Title: {result['title']}\n   Snippet: {result['snippet']}\n   URL: {result['link']}"
            for i, result in enumerate(search_results, start=1)
        ]
        logger.info(f"Websearch completed, retrieved {len(formatted_results)} results")
        return "Top Search Results:\n\n" + "\n\n".join(formatted_results)
    except Exception as e:
        logger.error(f"Error during web search: {str(e)}")
        return f"An error occurred during the web search: {str(e)}"
