import streamlit as st
import os

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from langchain_groq import ChatGroq # Changed to import ChatGroq
from langchain.prompts import ChatPromptTemplate

# Import the new dynamic template components
from htmlTemplates import css, user_template, get_bot_template, BOT_AVATARS_BASE64


# ------------------- Dynamic Bot Names -------------------
BOT_NAMES = {
    "Khamenei": "حاجاقا محسنی",
    "Sistani": "حاجاقا جوادی",
    "Golpaygani": "حاجاقا محمدی",
}


# ------------------- Conversation Chain Function -------------------
def get_conversation_chain(vectorstore, llm):
    memory = ConversationBufferMemory(
        memory_key='chat_history', return_messages=True
    )
    return ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vectorstore.as_retriever(),
        memory=memory,
    )


# ------------------- Classify Question -------------------
def classify_question(question, llm):
    prompt = ChatPromptTemplate.from_template("""
    متن زیر را بررسی کن و فقط یکی از این سه برچسب را برگردان:
    - greeting: اگر متن سلام و احوال‌پرسی یا جملات دوستانه است
    - islamic_fiqh: اگر پرسش درباره احکام شرعی و فقهی اسلام است
    - irrelevant: اگر غیر از این دو بود

    متن: {question}
    برچسب:
    """)
    chain = prompt | llm
    result = chain.invoke({"question": question})
    return result.content.strip().lower()


# ------------------- Generate Friendly Response -------------------
def generate_friendly_response(question, label, llm, selected_source_english):
    bot_name = BOT_NAMES.get(selected_source_english, "پاسخگوی شما")

    if "greeting" in label:
        prompt = ChatPromptTemplate.from_template(f"""
        کاربر به شما سلام یا احوال‌پرسی کرده است.
        نام شما {bot_name} هست.
        متن کاربر: {{question}}
        """)
    else:  # irrelevant
        prompt = ChatPromptTemplate.from_template(f"""
        کاربر پرسشی غیرمرتبط با احکام شرعی پرسیده است.
        تو به عنوان {bot_name}
        مودبانه و دوستانه به او بگویید که وظیفه‌ی شما فقط پاسخ به پرسش‌های شرعی است
        و او را به مطرح کردن یک سؤال دینی تشویق کنید.
        متن کاربر: {{question}}
        """)

    chain = prompt | llm
    result = chain.invoke({"question": question})
    return result.content.strip()


# ------------------- Handle User Input (Logic Only) -------------------
def handle_userinput(user_question):
    if st.session_state.conversation is None:
        st.warning("منابع هنوز آماده نشده‌اند.")
        return

    st.session_state.messages.append({"role": "user", "content": user_question})

    label = classify_question(user_question, st.session_state.llm)
    if "islamic_fiqh" in label:
        response = st.session_state.conversation({'question': user_question})
        bot_reply = response['answer']
    else:
        bot_reply = generate_friendly_response(user_question, label, st.session_state.llm, st.session_state.selected_source_english)

    st.session_state.messages.append({"role": "bot", "content": bot_reply})


# ------------------- Main Streamlit App -------------------
def main():
    st.set_page_config(page_title="شیخ جی پی تی", page_icon=":mosque:")
    st.write(css, unsafe_allow_html=True)

    st.markdown(
        "<h2 style=\"text-align: center;font-size: clamp(24px, 5vw, 48px); white-space: nowrap;\">پاسخگوی احکام شرعی 🕋</h2>",
        unsafe_allow_html=True
    )

    if "conversation" not in st.session_state:
        st.session_state.conversation = None
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "llm" not in st.session_state:
        # Changed LLM to use Groq API with ChatGroq
        st.session_state.llm = ChatGroq(
            model_name="llama-3.1-8b-instant",
            temperature=0.2,
            api_key=os.environ["GROQ_API_KEY"] # Changed to use GROQ_API_KEY
        )
    if "selected_source_english" not in st.session_state:
        st.session_state.selected_source_english = "Khamenei"

    current_bot_avatar_base64 = BOT_AVATARS_BASE64.get(st.session_state.selected_source_english)
    bot_template_with_avatar = get_bot_template(current_bot_avatar_base64)
    for msg in st.session_state.messages:
        if msg["role"] == "user":
            st.write(user_template.replace("{{MSG}}", msg["content"]), unsafe_allow_html=True)
        else:
            st.write(bot_template_with_avatar.replace("{{MSG}}", msg["content"]), unsafe_allow_html=True)

    if user_question := st.chat_input("سؤال شرعی خود را اینجا مطرح کنید..."):
        handle_userinput(user_question)
        st.rerun()

    with st.sidebar:
        st.subheader("منابع فقهی")

        SOURCE_MAPPINGS = {
            "آیت الله خامنه‌ای": "Khamenei",
            "آیت الله سیستانی": "Sistani",
            "آیت الله گلپایگانی": "Golpaygani"
        }
        persian_options = list(SOURCE_MAPPINGS.keys())
        current_persian_name = next(
            (key for key, value in SOURCE_MAPPINGS.items()
             if value == st.session_state.selected_source_english),
            persian_options[0]
        )

        selected_source_persian = st.selectbox(
            "مرجع تقلید خود را انتخاب کنید:",
            persian_options,
            index=persian_options.index(current_persian_name)
        )
        selected_source_english = SOURCE_MAPPINGS[selected_source_persian]

        if st.session_state.selected_source_english != selected_source_english:
            st.session_state.selected_source_english = selected_source_english
            st.session_state.conversation = None
            st.session_state.messages = []
            st.rerun()

        if st.session_state.conversation is None:
            placeholder = st.empty()
            placeholder.info(f"⏳ در حال بارگذاری منابع {selected_source_persian}...")
            try:
                embeddings = HuggingFaceEmbeddings(
                    model_name="heydariAI/persian-embeddings",
                    model_kwargs={'trust_remote_code': True},
                    cache_folder="/tmp/hf_cache"
                )
                vector_path = f"Resources/{st.session_state.selected_source_english}/faiss_index"
                vectorstore = FAISS.load_local(
                    vector_path,
                    embeddings,
                    allow_dangerous_deserialization=True
                )
                st.session_state.conversation = get_conversation_chain(vectorstore, st.session_state.llm)
                placeholder.success(
                    f"✅ منابع {selected_source_persian} آماده شدند. اکنون می‌توانید سؤال بپرسید.")
            except Exception as e:
                placeholder.error(f"⚠️ خطا در بارگذاری منابع: {e}")
        else:
            st.success(f"✅ منابع {selected_source_persian} بارگذاری شده‌اند")


if __name__ == '__main__':
    main()