import os
import streamlit as st
import logging
from langchain_groq import ChatGroq
from langchain_community.tools import ArxivQueryRun, WikipediaQueryRun, DuckDuckGoSearchRun
from langchain_community.utilities import ArxivAPIWrapper, WikipediaAPIWrapper
from langchain import hub
from langchain.agents import create_openai_tools_agent, AgentExecutor
from langchain.prompts import PromptTemplate
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.tools import Tool
from pydantic import BaseModel, Field
from typing import List, Dict

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logging.info("Streamlit app started")

# Load the API key
groq_api_key = "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" #os.getenv("GROQ_API_KEY")
if not groq_api_key:
    logging.error("GROQ API key is missing!")

# Define a query condensing template
condense_prompt_template = PromptTemplate(
    input_variables=["original_question", "conversation_history"],
    template="""
        Given the conversation history below, condense the user's query into a clear and specific question.
        Conversation History:
        {conversation_history}
        Original Question:
        {original_question}
        Condensed Question:"""
)

def condense_query(llm_model, original_question, conversation_history):
    prompt = condense_prompt_template.format(
        original_question=original_question,
        conversation_history=conversation_history
    )
    resp = llm_model.predict(prompt)
    return resp.strip()

# Initialize HuggingFace embeddings and FAISS vectorstore
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
vector_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
retriever_tool = vector_db.as_retriever()

# Define input/output models for OpenAI function-calling compliance
class RetrieveDocumentsInput(BaseModel):
    query: str = Field(..., description="Query to retrieve relevant documents.")

class RetrieveDocumentsOutput(BaseModel):
    documents: List[Dict[str, str]] = Field(..., description="List of retrieved documents, each containing 'content'.")

# Wrapper function for FAISS retriever
def retrieve_documents(query: str) -> List[Dict[str, str]]:
    results = retriever_tool.get_relevant_documents(query)
    return [{"content": doc.page_content} for doc in results]

# Define the FAISS tool
faiss_tool = Tool(
    name="retrieve_documents",
    description="Retrieve documents from FAISS vectorstore.",
    func=retrieve_documents,
)

# Initialize tools
arxiv_wrapper = ArxivAPIWrapper(top_k_results=1, doc_content_chars_max=250)
arxiv_tool = ArxivQueryRun(api_wrapper=arxiv_wrapper)

wiki_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=250)
wiki_tool = WikipediaQueryRun(api_wrapper=wiki_wrapper)

search_tool = DuckDuckGoSearchRun(name="Search")
tools = [arxiv_tool, wiki_tool, search_tool, faiss_tool]

# Set up the agent
prompt = hub.pull("hwchase17/openai-functions-agent")
llm = ChatGroq(model="mixtral-8x7b-32768", api_key=groq_api_key, streaming=True)
agent = create_openai_tools_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# Initialize session state for messages
if 'messages' not in st.session_state:
    st.session_state['messages'] = [{"role": "assistant", "content": "How can I help you?"}]

# Display chat messages
for msg in st.session_state['messages']:
    st.chat_message(msg["role"]).write(msg["content"])

# Extract conversation history
conversation_history = "\n".join(
    [f"{msg['role']}: {msg['content']}" for msg in st.session_state['messages']]
)

# CSS for input box at bottom of screen
st.markdown("""
    <style>
    .fixed-bottom-input-container {
        position: fixed;
        bottom: 0;
        width: 100%;
        background-color: white;
        padding: 10px 0;
        border-top: 1px solid #ddd;
    }
    .fixed-bottom-input {
        width: 100%;
        padding: 10px;
        font-size: 16px;
    }
    </style>
""", unsafe_allow_html=True)

# Chat input and processing
user_input = st.text_input("Type your message here...", key="user_input", label_visibility="collapsed")

# If there's input, process it
if user_input:
    condensed_question = condense_query(llm, user_input, conversation_history)
    
    # Placeholder for streaming the assistant's response
    response_placeholder = st.empty()
    response_text = ""
    
    # Stream response tokens
    for token in agent_executor.stream({"input": condensed_question}):
        # Ensure the token contains the "output" key
        if "output" in token:
            response_text += token["output"]
            response_placeholder.write(response_text)
        else:
            logging.warning("Received token without 'output' key: %s", token)
    
    # Continue with the chat history and message updates
    st.session_state.messages.append({"role": "user", "content": user_input})
    st.session_state.messages.append({"role": "assistant", "content": response_text})
    
    # Display the user input and response
    st.chat_message("user").write(user_input)
    st.chat_message("assistant").write(response_text)
    
    # Reset the input box content
    user_input = ""