import os
from langchain.chains import RetrievalQA
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import FAISS
import gradio as gr

# Load secrets
openai_api_key = os.environ.get("OPENAI_API_KEY")
app_password = os.environ.get("APP_PASSWORD")
app_username = os.environ.get("APP_USERNAME", "user")

if not openai_api_key:
    raise ValueError("OPENAI_API_KEY is not set.")
if not app_password:
    raise ValueError("APP_PASSWORD is not set.")

# Load vectorstore
def load_vectorstore():
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    return FAISS.load_local(".", embeddings, allow_dangerous_deserialization=True)

vectorstore = load_vectorstore()

# Load LLM
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.1)

# Setup QA chain
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectorstore.as_retriever(),
    return_source_documents=True
)
def chat_with_pdf(query):
    if not query.strip():
        return "❗ Please enter a question."
    
    result = qa_chain.invoke({"query": query})
    answer = result["result"]
    sources = result.get("source_documents", [])

    if sources:
        formatted_sources = "".join(
            f"\n\n📄 Source {i+1}:\n{doc.page_content.strip()}"
            for i, doc in enumerate(sources)
        )
        return f"🧠 **Answer:**\n{answer}{formatted_sources}"
    else:
        return f"🧠 **Answer:**\n{answer}\n\nℹ️ No source documents found."


# App state
authenticated_users = set()

def login(username, password):
    if username == app_username and password == app_password:
        authenticated_users.add(username)
        return gr.update(visible=True), gr.update(visible=False), "", ""
    else:
        return gr.update(visible=False), gr.update(visible=True), "", "❌ Invalid credentials."

with gr.Blocks() as demo:
    gr.Markdown("# 🔐 Secure  Chatbot")

    with gr.Group(visible=True) as login_box:
        gr.Markdown("### Login to access the chatbot")
        login_user = gr.Textbox(label="Username")
        login_pass = gr.Textbox(label="Password", type="password")
        login_button = gr.Button("Login")
        login_message = gr.Markdown("")

    with gr.Group(visible=False) as chat_box:
        gr.Markdown("## 📚 Ask questions ")
        query_input = gr.Textbox(label="Question", lines=2, placeholder="Type your question here...")
        submit_button = gr.Button("Submit")
        output = gr.Markdown()

    login_button.click(
        fn=login,
        inputs=[login_user, login_pass],
        outputs=[chat_box, login_box, login_user, login_message]
    )
    submit_button.click(fn=chat_with_pdf, inputs=query_input, outputs=output)

if __name__ == "__main__":
    demo.launch(share=True)