import streamlit as st
import os
import base64

# Title of the page
st.title("Physics Test Sets")

# Create a column layout for displaying links and PDFs
col1, col2, col3 = st.columns(3)

# Function to display a clickable download link for the PDF
def create_pdf_link(pdf_name, file_path):
    # Ensure the file name ends with '.pdf'
    if not pdf_name.endswith(".pdf"):
        pdf_name += ".pdf"
    
    if os.path.exists(file_path):
        with open(file_path, "rb") as f:
            st.download_button(label=pdf_name, data=f, file_name=pdf_name, mime="application/pdf")
    else:
        st.error(f"File {pdf_name} not found!")

# Function to display the PDF in the Streamlit page
def display_pdf(pdf_file):
    with open(pdf_file, "rb") as f:
        base64_pdf = base64.b64encode(f.read()).decode('utf-8')
    pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="700" height="900" type="application/pdf"></iframe>'
    st.markdown(pdf_display, unsafe_allow_html=True)

# Display the first PDF link as a clickable download button in the first column
with col1:
    pdf_path_1 = os.path.join("psets", "eJackson_set1.pdf")
    create_pdf_link("Electrodynamics Test Set 1", pdf_path_1)
    if st.button("View Electrodynamics Test Set 1"):
        display_pdf(pdf_path_1)

with col2:
    pdf_path_2 = os.path.join("psets", "Griffiths_Example.pdf")
    create_pdf_link("Griffiths Example", pdf_path_2)
    if st.button("View Griffiths Example"):
        display_pdf(pdf_path_2)

# Add more columns and rows as needed for additional PDFs