import streamlit as st
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.schema import Document
from transformers import pipeline
import pickle
import time
from string import Template


def load_faiss_index(index_file="faiss_course_index.pkl"):
    with open(index_file, "rb") as f:
        return pickle.load(f)

def get_faiss_retriever(index_file="faiss_course_index.pkl"):
    faiss_index = load_faiss_index(index_file)
    documents = [Document(page_content=item['description'], metadata=item) for item in faiss_index['metadata']]
    retriever = FAISS.from_documents(documents, HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")).as_retriever()
    return retriever

def get_answer_from_t5(query, context):
    t5_pipeline = pipeline("text2text-generation", model="t5-small", device=0)
    input_text = f"Question: {query} Context: {context}"
    answer = t5_pipeline(input_text, max_length = 500,truncation = True)
    return answer[0]['generated_text']

    
def search_and_answer(query, index_file="faiss_course_index.pkl"):
    retriever = get_faiss_retriever(index_file=index_file)
    retrieved_docs = retriever.get_relevant_documents(query)

    if not retrieved_docs:
        return "No relevant courses found."

    count = len(retrieved_docs)
    print(f"Related Courses: {count}")
    course_info = [
        (doc.metadata.get('title', 'No Title'), doc.metadata.get('lessons', 'No lessons listed'))
        for doc in retrieved_docs
    ]
    response = "\nRelated Courses:\n"
    for course_title, course_lessons in course_info:
        response += f"Course Title: {course_title}\nCourse Lessons: {course_lessons}\n\n"

    context = retrieved_docs[0].page_content
    summary = get_answer_from_t5(query, context)
    response += f"Summary: {summary}"

    return response, course_info, summary

def load_card_template():
    with open("card_template.html",'r') as file:
        return file.read()

def main():
    st.title("📚 Course Search WebApp")
    query = st.text_input("🔍 Enter your query:")

    if st.button("Search"):
        with st.spinner("Fetching results..."):
            start_time = time.time()
            response, course_info, summary = search_and_answer(query, index_file="faiss_course_index.pkl")

        if response == "No relevant courses found.":
            st.error(response)
        else:
            st.success("✅ Results fetched successfully!")
            st.write(f"⏱ **Time Taken:** {time.time() - start_time:.2f} seconds")

            # Load the card template HTML
            card_template = load_card_template()

            # Display related courses as cards in two columns
            st.subheader("🎓 Related Courses:")
            st.write("<div style='display: flex; flex-wrap: wrap; justify-content: space-between;'>", unsafe_allow_html=True)
            for course_title, course_lessons in course_info:
                card_html = Template(card_template).substitute(course_title=course_title, course_lessons=course_lessons)
                st.write(card_html, unsafe_allow_html=True)
            st.write("</div>", unsafe_allow_html=True)

            # Display summary with better layout
            st.subheader("📝 Summary:")
            st.info(summary[:500])

    st.markdown("<p style='text-align: center;'>Developed with ❤️ by Prachi</p>", unsafe_allow_html=True)

if __name__ == "__main__":
    main()