import os
from fastapi import FastAPI, File, UploadFile
from io import BytesIO
from PyPDF2 import PdfReader, PdfWriter
from fastapi.responses import StreamingResponse

app = FastAPI()

def compress_pdf(input_pdf: BytesIO) -> BytesIO:
    reader = PdfReader(input_pdf)
    writer = PdfWriter()
    
    for page_num in range(len(reader.pages)):
        writer.add_page(reader.pages[page_num])
    
    output_pdf = BytesIO()
    writer.write(output_pdf)
    output_pdf.seek(0)
    return output_pdf

@app.post("/compress-pdf/")
async def compress_pdf_endpoint(file: UploadFile = File(...)):
    input_pdf = BytesIO(await file.read())
    compressed_pdf = compress_pdf(input_pdf)
    return StreamingResponse(compressed_pdf, media_type="application/pdf", headers={"Content-Disposition": f"attachment; filename={file.filename}"})