import time
import os
import json
import random
import streamlit as st
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from vectorize_documents import embeddings
from deep_translator import GoogleTranslator
from googlesearch import search
from datetime import datetime

# Set up working directory and API configuration
working_dir = os.path.dirname(os.path.abspath(__file__))
config_data = json.load(open(f"{working_dir}/config.json"))
os.environ["GROQ_API_KEY"] = config_data["GROQ_API_KEY"]

def setup_vectorstore():
    persist_directory = f"{working_dir}/vector_db_dir"
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embeddings
    )
    return vectorstore

def chat_chain(vectorstore):
    from langchain_groq import ChatGroq

    llm = ChatGroq(
        model="llama-3.1-70b-versatile",
        temperature=0
    )
    retriever = vectorstore.as_retriever()
    memory = ConversationBufferMemory(
        llm=llm,
        output_key="answer",
        memory_key="chat_history",
        return_messages=True
    )

    chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        chain_type="stuff",
        memory=memory,
        verbose=True,
        return_source_documents=True
    )
    return chain

def fetch_daily_quote():
    """Fetch a daily Bhagavad Gita quote with its URL."""
    query = "Bhagavad Gita inspirational quotes"
    results = list(search(query, num_results=5))  # Convert generator to list
    if results:
        chosen_result = random.choice(results)
        return chosen_result
    return "Explore the Bhagavad Gita and Yoga Sutras for timeless wisdom!"

def get_daily_quote():
    """Get or refresh the daily quote based on the date."""
    current_date = datetime.now().date()
    if "daily_quote_date" not in st.session_state or st.session_state.daily_quote_date != current_date:
        # Fetch new quote and save the current date
        st.session_state.daily_quote_date = current_date
        st.session_state.daily_quote = fetch_daily_quote()

    return st.session_state.daily_quote

# Streamlit UI
st.set_page_config(
    page_title="Bhagavad Gita & Yoga Sutras Assistant",
    page_icon="🕉️",
    layout="wide"
)

st.markdown(
    """
    <div style="text-align: center;">
        <h1 style="color: #4CAF50;">Wisdom Query Assistant</h1>
        <p style="font-size: 18px;">Explore timeless wisdom with the guidance of a knowledgeable assistant.</p>
    </div>
    """,
    unsafe_allow_html=True
)

# User name functionality
if "user_name" not in st.session_state:
    st.session_state.user_name = ""

if "chat_started" not in st.session_state:
    st.session_state.chat_started = False

if not st.session_state.chat_started:
    st.markdown("<h3 style='text-align: center;'>Welcome! Before we begin, please enter your name:</h3>", unsafe_allow_html=True)
    user_name = st.text_input("Enter your name:", placeholder="Your Name", key="name_input")
    start_button = st.button("Start Chat")

    if start_button and user_name.strip():
        st.session_state.user_name = user_name.strip()
        st.session_state.chat_started = True
        st.success(f"Hello {st.session_state.user_name}! How can I assist you today?")

# Display the daily quote
daily_quote = get_daily_quote()
st.markdown(
    f"""
    <div style="text-align: center; background-color: #f0f8ff; padding: 10px; border-radius: 5px; margin-bottom: 20px;">
        <h4>🌟 Daily Wisdom: <a href="{daily_quote}" target="_blank">{daily_quote}</a></h4>
    </div>
    """,
    unsafe_allow_html=True
)

if st.session_state.chat_started:
    # Set up vectorstore and chat chain
    vectorstore = setup_vectorstore()
    chain = chat_chain(vectorstore)

    # Select language
    selected_language = st.selectbox("Select your preferred language:", options=[
        "English", "Hindi", "Bengali", "Telugu", "Marathi", "Tamil", "Urdu", "Gujarati", "Malayalam", "Kannada",
        "Punjabi", "Odia", "Maithili", "Sanskrit", "Santali", "Kashmiri", "Nepali", "Dogri", "Manipuri", "Bodo",
        "Sindhi", "Assamese", "Konkani", "Awadhi", "Rajasthani", "Haryanvi", "Bihari", "Chhattisgarhi", "Magahi"
    ], index=0)

    # Display chat history
    st.markdown("### 💬 Chat History")
    if "chat_history" in st.session_state:
        for chat in st.session_state.chat_history:
            st.markdown(f"**{st.session_state.user_name}:** {chat['question']}")
            st.markdown(f"**Assistant:** {chat['answer']}")
            st.markdown("---")

    # Input box for new query
    st.markdown(f"### Ask a new question, {st.session_state.user_name}:")
    with st.form("query_form", clear_on_submit=True):
        user_query = st.text_input("Your question:", key="query_input", placeholder="Type your query here...")
        submitted = st.form_submit_button("Submit")

    if submitted and user_query.strip():
        start_time = time.time()
        response = chain({"question": user_query.strip()})
        end_time = time.time()

        answer = response.get("answer", "No answer found.")
        source_documents = response.get("source_documents", [])
        execution_time = round(end_time - start_time, 2)

        # Translate response if needed
        if selected_language != "English":
            translator = GoogleTranslator(source="en", target=selected_language.lower())
            translated_answer = translator.translate(answer)
        else:
            translated_answer = answer

        # Save chat history
        if "chat_history" not in st.session_state:
            st.session_state.chat_history = []
        st.session_state.chat_history.append({
            "question": user_query.strip(),
            "answer": translated_answer
        })

        # Display source documents if available
        if source_documents:
            with st.expander("📜 Source Documents"):
                for i, doc in enumerate(source_documents):
                    st.write(f"**Document {i + 1}:** {doc.page_content}")

        st.write(f"**🧙‍♂️ Assistant:** {translated_answer}")
        st.write(f"_Response time: {execution_time} seconds_")

    # Sharing options
    st.markdown(
        """
        <div style="text-align: center;">
            <a href="https://wa.me/?text=Explore%20the%20Bhagavad%20Gita%20%26%20Yoga%20Sutras%20Assistant!%20Check%20it%20out%20here:%20https://your-platform-link" target="_blank">
                <img src="https://img.icons8.com/color/48/whatsapp.png" alt="WhatsApp" style="margin-right: 10px;">
            </a>
            <a href="https://www.linkedin.com/shareArticle?mini=true&url=https://your-platform-link&title=Explore%20Wisdom%20with%20Our%20Assistant" target="_blank">
                <img src="https://img.icons8.com/color/48/linkedin.png" alt="LinkedIn">
            </a>
        </div>
        """,
        unsafe_allow_html=True
    )