# Standard Library Imports
import logging
import os

# Third-party Imports
from dotenv import load_dotenv
import chromadb
import gradio as gr
from huggingface_hub import snapshot_download

# LlamaIndex (Formerly GPT Index) Imports
from llama_index.core import VectorStoreIndex, SimpleKeywordTableIndex
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core.llms import MessageRole
from llama_index.core.memory import ChatSummaryMemoryBuffer
from llama_index.core.tools import RetrieverTool, ToolMetadata, FunctionTool
from llama_index.agent.openai import OpenAIAgent
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.llms.perplexity import Perplexity
from llama_index.core import Settings, QueryBundle
from llama_index.core.schema import NodeWithScore
from llama_index.core.retrievers import (
    BaseRetriever,
    VectorIndexRetriever,
    KeywordTableSimpleRetriever,
)
from typing import List
from llama_index.core import get_response_synthesizer
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.postprocessor.cohere_rerank import CohereRerank

_cached_indices = {}
_cached_tools = None

class HybridRetriever(BaseRetriever):
    """Hybrid retriever that performs both semantic search and keyword search."""

    def __init__(
        self,
        vector_retriever: VectorIndexRetriever,
        keyword_retriever: KeywordTableSimpleRetriever,
        max_retrieve: int = 10,
    ) -> None:
        """Init params."""

        self._vector_retriever = vector_retriever
        self._keyword_retriever = keyword_retriever
        self._max_retrieve = max_retrieve
        super().__init__()

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """Retrieve nodes given query."""

        vector_nodes = self._vector_retriever.retrieve(query_bundle)
        keyword_nodes = self._keyword_retriever.retrieve(query_bundle)

        resulting_nodes = []
        node_ids_added = set()
        for i in range(min(len(vector_nodes), len(keyword_nodes))):
            vector_node = vector_nodes[i]
            if vector_node.node.node_id not in node_ids_added:
                resulting_nodes += [vector_node]
                node_ids_added.add(vector_node.node.node_id)

            keyword_node = keyword_nodes[i]
            if keyword_node.node.node_id not in node_ids_added:
                resulting_nodes += [keyword_node]
                node_ids_added.add(keyword_node.node.node_id)

        return resulting_nodes

def retrieve_all_nodes_from_vector_index(vector_index, query="Whatever", similarity_top_k=100000000):
    # Set similarity_top_k to a large number to retrieve all the nodes
    vector_retriever = vector_index.as_retriever(similarity_top_k=similarity_top_k)

    # Retrieve all nodes
    all_nodes = vector_retriever.retrieve(query)
    nodes = [item.node for item in all_nodes]

    return nodes


def web_search(query: str) -> str:
    """
    Search the web for current information using Perplexity API.
    
    Args:
        query (str): The search query to look up on the web
        
    Returns:
        str: Search results from the web
    """
    try:
        perplexity_api_key = os.getenv("PERPLEXITY_API_KEY")
        if not perplexity_api_key:
            return "Error: Perplexity API key not found. Please provide your Perplexity API key."
        
        perplexity_llm = Perplexity(
            api_key=perplexity_api_key,
            model="sonar",
            temperature=0.2
        )

        search_prompt = f"Search the web for current information about: {query}. Provide a comprehensive and accurate response based on recent sources."

        logging.info(f"Performing web search: {search_prompt}")

        response = perplexity_llm.complete(search_prompt)
        
        return str(response)
        
    except Exception as e:
        logging.error(f"Error in web search: {e}")
        return f"Error performing web search: {str(e)}"


load_dotenv()

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
logging.getLogger("httpx").setLevel(logging.WARNING)

PROMPT_SYSTEM_MESSAGE = """You are an AI teacher, answering questions from students of an applied AI course on Large Language Models (LLMs or llm) and Retrieval Augmented Generation (RAG) for LLMs. 
Topics covered include training models, fine-tuning models, giving memory to LLMs, prompting tips, hallucinations and bias, vector databases, transformer architectures, embeddings, RAG frameworks such as 
Langchain and LlamaIndex, making LLMs interact with tools, AI agents, reinforcement learning with human feedback (RLHF). Questions should be understood in this context. Your answers are aimed to teach 
students, so they should be complete, clear, and easy to understand. Use the available tools to gather insights pertinent to the field of AI.

You have access to tools that may include:
1. "AI_Information_related_resources" - Use this tool to search the local knowledge base for comprehensive information about AI concepts, frameworks, and methodologies.
2. "web_search" - If available, use this tool to search the web for current, up-to-date information about recent developments, new tools, or information that might not be in the local knowledge base.

Choose the appropriate tool based on the question:
- For fundamental AI concepts, frameworks, and established methodologies, use the local knowledge base tool.
- For recent developments, new tools, current events in AI, or information that might be too recent for the local knowledge base, use the web search tool if available.
- You can use both tools if needed to provide comprehensive answers.

Only some information returned by the tools might be relevant to the question, so ignore the irrelevant part and answer the question with what you have. Your responses are exclusively based on the output provided 
by the tools. Refrain from incorporating information not directly obtained from the tool's responses.
If a user requests further elaboration on a specific aspect of a previously discussed topic, you should reformulate your input to the tool to capture this new angle or more profound layer of inquiry. Provide 
comprehensive answers, ideally structured in multiple paragraphs, drawing from the tool's variety of relevant details. The depth and breadth of your responses should align with the scope and specificity of the information retrieved. 
Should the tool response lack information on the queried topic, politely inform the user that the question transcends the bounds of your current knowledge base, citing the absence of relevant content in the tool's documentation. 
At the end of your answers, always invite the students to ask deeper questions about the topic if they have any.
Do not refer to the documentation directly, but use the information provided within it to answer questions. If code is provided in the information, share it with the students. It's important to provide complete code blocks so 
they can execute the code when they copy and paste them. Make sure to format your answers in Markdown format, including code blocks and snippets.
"""

TEXT_QA_TEMPLATE = """
You must answer only related to AI, ML, Deep Learning and related concepts queries.
Always leverage the retrieved documents to answer the questions, don't answer them on your own.
If the query is not relevant to AI, say that you don't know the answer.
"""


def download_knowledge_base_if_not_exists():
    """Download the knowledge base from the Hugging Face Hub if it doesn't exist locally"""
    if not os.path.exists("data/ai_tutor_knowledge"):
        os.makedirs("data/ai_tutor_knowledge")

        logging.warning(
            f"Vector database does not exist at 'data/', downloading from Hugging Face Hub..."
        )
        snapshot_download(
            repo_id="jaiganesan/ai_tutor_knowledge_vector_Store",
            local_dir="data/ai_tutor_knowledge",
            repo_type="dataset",
        )
        logging.info(f"Downloaded vector database to 'data/ai_tutor_knowledge'")


def clear_cache():
    """Clear the cached indices and tools to force recreation"""
    global _cached_indices, _cached_tools
    _cached_indices = {}
    _cached_tools = None
    logging.info("Cleared cached indices and tools")


def get_tools(db_collection="ai_tutor_knowledge"):
    global _cached_indices, _cached_tools
    
    if _cached_tools is not None:
        logging.info("Using cached tools")
        return _cached_tools
    
    if db_collection not in _cached_indices:
        logging.info(f"Creating indices for collection: {db_collection}")
        
        db = chromadb.PersistentClient(path=f"data/{db_collection}")
        chroma_collection = db.get_or_create_collection(db_collection)
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

        index = VectorStoreIndex.from_vector_store(
            vector_store=vector_store,
            show_progress=True,
            use_async=True,
            embed_model=Settings.embed_model,
            verbose=True,
        )
        nodes = retrieve_all_nodes_from_vector_index(index)
        keyword_index = SimpleKeywordTableIndex(nodes=nodes)
        
        _cached_indices[db_collection] = {
            'vector_index': index,
            'keyword_index': keyword_index,
            'nodes': nodes
        }
        logging.info(f"Cached indices for collection: {db_collection}")
    else:
        logging.info(f"Using cached indices for collection: {db_collection}")
    
    cached_data = _cached_indices[db_collection]
    index = cached_data['vector_index']
    keyword_index = cached_data['keyword_index']
    
    vector_retriever = VectorIndexRetriever(
        index=index,
        similarity_top_k=15,
        embed_model=Settings.embed_model,
        use_async=True,
        verbose=True
    )

    keyword_retriever = KeywordTableSimpleRetriever(index=keyword_index, num_chunks_per_query=6)
    hybrid_retriever = HybridRetriever(vector_retriever, keyword_retriever, max_retrieve=6)
    
    cohere_api_key = os.getenv("CO_API_KEY")
    if cohere_api_key:
        cohere_rerank = CohereRerank(top_n=4, api_key=cohere_api_key, model='rerank-english-v3.0')
        hybrid_retriever.node_postprocessors = [cohere_rerank]
        logging.info("Cohere reranking enabled")
    else:
        logging.warning("CO_API_KEY not found. Reranking disabled.")
    
    tools = [
        RetrieverTool(
            retriever=hybrid_retriever,
            metadata=ToolMetadata(
                name="AI_Information_related_resources",
                description="Useful for info related to artificial intelligence, ML, deep learning. It gathers the info from local data.",
            ),
        )
    ]
    
    perplexity_api_key = os.getenv("PERPLEXITY_API_KEY")
    if perplexity_api_key:
        tools.append(
            FunctionTool.from_defaults(
                fn=web_search,
                name="web_search",
                description="Search the web for current and up-to-date information about AI, ML, deep learning, and related topics. Use this tool when you need the latest information that might not be in the local knowledge base.",
            )
        )
        logging.info("Perplexity web search tool enabled")
    else:
        logging.info("PERPLEXITY_API_KEY not found. Web search tool disabled.")
    
    _cached_tools = tools
    logging.info("Cached tools for reuse")
    
    return tools


def generate_completion(query, history, memory, openai_key, cohere_key, perplexity_key):
    logging.info(f"User query: {query}")

    try:
        chat_list = memory.get()
        if len(chat_list) != 0:
            user_index = [i for i, msg in enumerate(chat_list) if msg.role == MessageRole.USER]
            if len(user_index) > len(history):
                user_index_to_remove = user_index[len(history)]
                chat_list = chat_list[:user_index_to_remove]
                memory.set(chat_list)
        logging.info(f"chat_history: {len(memory.get())} {memory.get()}")
        logging.info(f"gradio_history: {len(history)} {history}")

        # Create agent
        tools = get_tools(db_collection="ai_tutor_knowledge")
        agent = OpenAIAgent.from_tools(
            llm=Settings.llm,
            memory=memory,
            tools=tools,
            system_prompt=PROMPT_SYSTEM_MESSAGE,
        )

        completion = agent.stream_chat(query)
        answer_str = ""
        
        for token in completion.response_gen:
            answer_str += token
            yield answer_str
            
    except Exception as e:
        logging.error(f"Error in generate_completion: {e}")
        error_message = f"Sorry, I encountered an error while processing your request: {str(e)}"
        yield error_message



def launch_ui():
    with gr.Blocks(
        fill_height=True,
        title="AI Tutor 🤖",
        analytics_enabled=True,
        theme=gr.themes.Soft(),
    ) as demo:

        memory_state = gr.State(
            lambda: ChatSummaryMemoryBuffer.from_defaults(
                token_limit=120000,
            )
        )

        with gr.Row():
            with gr.Column(scale=1):
                open_ai_api_key = gr.Textbox(
                    label="OpenAI API Key", 
                    type="password",
                    placeholder="Enter your OpenAI API key..."
                )
            with gr.Column(scale=1):
                cohere_api_key = gr.Textbox(
                    label="Cohere API Key (Optional - for reranking)", 
                    type="password",
                    placeholder="Enter your Cohere API key..."
                )
            with gr.Column(scale=1):
                perplexity_api_key = gr.Textbox(
                    label="Perplexity API Key (Optional - for web search)", 
                    type="password",
                    placeholder="Enter your Perplexity API key..."
                )
        
        def init_clients(openai_key, cohere_key, perplexity_key):
            clear_cache()
            
            if openai_key:
                Settings.llm = OpenAI(temperature=0, model="gpt-4o-mini", api_key=openai_key)
                Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small", api_key=openai_key)
                os.environ["OPENAI_API_KEY"] = openai_key
            if cohere_key:
                os.environ["CO_API_KEY"] = cohere_key
                logging.info("Cohere API key set for reranking")
            if perplexity_key:
                os.environ["PERPLEXITY_API_KEY"] = perplexity_key
                logging.info("Perplexity API key set for web search")
        
        open_ai_api_key.change(init_clients, inputs=[open_ai_api_key, cohere_api_key, perplexity_api_key])
        cohere_api_key.change(init_clients, inputs=[open_ai_api_key, cohere_api_key, perplexity_api_key])
        perplexity_api_key.change(init_clients, inputs=[open_ai_api_key, cohere_api_key, perplexity_api_key])


        chatbot = gr.Chatbot(
            scale=1,
            placeholder="<strong>AI Tutor 🤖: A Question-Answering Bot for anything AI-related</strong><br><br>Ask me anything about AI, ML, Deep Learning, and related topics!",
            show_label=False,
            show_copy_button=True,
            height=500,
            bubble_full_width=False,
        )

        gr.ChatInterface(
            fn=generate_completion,
            chatbot=chatbot,
            additional_inputs=[memory_state, open_ai_api_key, cohere_api_key, perplexity_api_key],
            retry_btn=None,
            undo_btn=None,
            clear_btn="Clear Chat",
            submit_btn="Send Message",
        )

        demo.queue(default_concurrency_limit=64)
        demo.launch(debug=True, share=False) # Set share=True to share the app online


if __name__ == "__main__":
    # Download the knowledge base if it doesn't exist
    download_knowledge_base_if_not_exists()

    # Set up llm and embedding model
    Settings.llm = OpenAI(temperature=0, model="gpt-4o-mini")
    Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

    # launch the UI
    launch_ui()