import streamlit as st
import PyPDF2
import os
from io import BytesIO

st.subheader("合并PDF")

pdf_merger = PyPDF2.PdfMerger()
file_count = 0
uploaded_files = st.file_uploader("选择PDF文件", type="pdf",accept_multiple_files=True)
for uploaded_file in uploaded_files:
    bytes_data = uploaded_file.read()
    pdf_file = BytesIO(bytes_data)
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    pdf_merger.append(pdf_reader)
    file_count += 1

if st.button("合并"):
    if file_count < 2:
        st.error(f"至少需要两个文件才能合并,当前文件数:{file_count}")
        st.stop()
    with st.spinner("合并中..."):
        output_path = os.path.join(os.getcwd(), 'merged.pdf')
        with open(output_path, 'wb') as output_file:
            pdf_merger.write(output_file)
        with open(output_path, 'rb') as file:
            st.download_button(label="下载合并后的 PDF", 
                            data=file, 
                            file_name='merged.pdf', 
                            mime='application/pdf')