import os
import gradio as gr
from langchain_community.document_loaders import TextLoader
from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpoint
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_huggingface import ChatHuggingFace # Add this to your imports

# 2026 Imports: These now live in the classic/retrieval modules
from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain

# 1. Setup API Token from Secrets
sec_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")

# 2. Load and Process Knowledge Base
# Make sure "knowledge.txt" exists in your Space files!
loader = TextLoader("knowledge.txt")
docs = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)

# 3. Create Vector Store
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
retriever = vectorstore.as_retriever()

# 4. Initialize LLM (The Fix: Task="conversational")
llm = HuggingFaceEndpoint(
    repo_id="Qwen/Qwen2.5-7B-Instruct",
    task="conversational",  # Matches the provider's requirement
    huggingfacehub_api_token=sec_token,
    temperature=0.1,
)

chat_model = ChatHuggingFace(llm=llm)

# 5. Define the RAG Prompt
system_prompt = (
    "You are a helpful assistant for Examfriendly tuition. "
    "Use the following pieces of retrieved context to answer the question. "
    "If you don't know the answer, say you don't know. "
    "\n\n"
    "{context}"
)

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        ("human", "{input}"),
    ]
)

# 6. Build the Chain
question_answer_chain = create_stuff_documents_chain(chat_model, prompt)
rag_chain = create_retrieval_chain(retriever, question_answer_chain)

# 7. Define the Chat Function for Gradio
def chat(message, history):
    try:
        # ChatHuggingFace expects "input" inside the dictionary
        result = rag_chain.invoke({"input": message})
        return result["answer"]
    except Exception as e:
        return f"System Error: {str(e)}"

# 8. Launch the Interface
demo = gr.ChatInterface(
    fn=chat,
    title="Examfriendly AI Tutor",
    description="Ask me anything about your tuition courses!",
)

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